# windows_usbkill
# Author: Willy

import wmi
import subprocess
import os
import pystray
import threading
import sys
import pystray
import PIL.Image
import ctypes
import logging
import pythoncom
import time
import tkinter as tk

from tkinter import ttk
from tkinter.ttk import Label
from win11toast import toast
from pystray import MenuItem as item

if sys.platform != 'win32':
    print("This script only works on Windows!")
    logging.info("[ERROR] This script only works on Windows!")
    os._exit(0)

# Fix Hi-DPI
ctypes.windll.shcore.SetProcessDpiAwareness(2)

# Logging
logging.basicConfig(filename="log.txt",
    format='%(asctime)s | %(message)s',
    level=logging.INFO)

# Load Image 512x512
current_dir_image = os.path.dirname(os.path.abspath(__file__))
image_file = os.path.join(current_dir_image, 'assets', 'icon_512x512.png')
try:
    image = PIL.Image.open(image_file)
except FileNotFoundError:
        asset_error = (None, 'File missing! Try to pull from Git!', 'Windows usbkill | Fatal Error!', 0x10)
        logging.info("[ERROR] File missing!")
        ctypes.windll.user32.MessageBoxW(*asset_error)
        os._exit(0)
# Load Image 128x128
current_dir_image = os.path.dirname(os.path.abspath(__file__))
image_file2 = os.path.join(current_dir_image, 'assets', 'icon_128x128.png')
try:
    image2 = PIL.Image.open(image_file2)
except FileNotFoundError:
        asset_error = (None, 'File missing! Try to pull from Git!', 'Windows usbkill | Fatal Error!', 0x10)
        logging.info("[ERROR] File missing!")
        ctypes.windll.user32.MessageBoxW(*asset_error)
        os._exit(0)
# Load ico
current_dir_image = os.path.dirname(os.path.abspath(__file__))
image_icon = os.path.join(current_dir_image, 'assets', 'icon.ico')
try:
    image3 = PIL.Image.open(image_icon)
except FileNotFoundError:
        asset_error = (None, 'File missing! Try to pull from Git!', 'Windows usbkill | Fatal Error!', 0x10)
        logging.info("[ERROR] File missing!")
        ctypes.windll.user32.MessageBoxW(*asset_error)
        os._exit(0)
# load exit-icon
current_dir_image = os.path.dirname(os.path.abspath(__file__))
exit_icon = os.path.join(current_dir_image, 'assets', 'exit-icon_512x512.png')
try:
    image4 = PIL.Image.open(exit_icon)
except FileNotFoundError:
        asset_error = (None, 'File missing! Try to pull from Git!', 'Windows usbkill | Fatal Error!', 0x10)
        logging.info("[ERROR] File missing!")
        ctypes.windll.user32.MessageBoxW(*asset_error)
        os._exit(0)

logging.info("[INFO] Script Started!")

# Variables
is_paused = False
event = threading.Event()

# Scan
def start_scan():
    pythoncom.CoInitialize()
    global is_paused
    while True:
        if is_paused == False:
            print("[Info] Scanning...")
            raw_wql = "SELECT * FROM __InstanceCreationEvent WITHIN 2 WHERE TargetInstance ISA \'Win32_USBHub\'"
            a = wmi.WMI ()
            watcher = a.watch_for(raw_wql=raw_wql)
            if is_paused == True:
                break
            while 1:
                usb = watcher ()
                if usb:
                    if is_paused == True:
                        break
                logging.info("[TRIGGERED] USB INSERTED! Shutting down!")
                print("[TRIGGERED] USB INSERTED! Shutting down!")
                os.system("shutdown /s /f /t 0")
        else:
            break

# Pause
def toggle_pause():
    pythoncom.CoInitialize()
    global is_paused
    is_paused = not is_paused
    if is_paused:
        print("[Info] Paused!")
    else:
        print("[Info] Resumed!")
        scan_thread = threading.Thread(target=start_scan)
        scan_thread.start()


# About Page
def about_section():
    root = tk.Tk()
    root.geometry('600x450')
    root.resizable(False, False)
    root.title('About | Windows usbkill')
    root.iconbitmap(image_icon)

        # Declare Exit
    def close_window():
        root.destroy()

    # Image

    photo = tk.PhotoImage(file=image_file2)
    image_label = ttk.Label(
        root,
        image=photo,
        padding=5
)
    image_label.pack()

    # Label 1
    label = ttk.Label(
        root,
        text='Windows usbkill',
        font=("Helvetica", 14, "bold"))

    label.pack(ipadx=10, ipady=10)

    label2 = ttk.Label(
        root,
        text = '''
        Author: Willy
        Version: 0.2 BETA (9/12/2023)
        License: GNU GENERAL PUBLIC LICENSE 3
        ''',
        font=("Helvetica", 10))

    # Label 2
    label2.pack(ipadx=200, ipady=20)

    # Exit Button
    button = tk.Button(root, text="Close Window", command=close_window)
    button.pack()

    # loop
    root.mainloop()

# Trigger about section
def trigger_about_section():
    about_thread = threading.Thread(target=about_section)
    about_thread.start()

# Exit
def terminate():
    pythoncom.CoInitialize()
    print("[Info] Exited!")
    icon = pystray.Icon
    icon.visible = False
    def exit_toast():
        toast(app_id="Windows usbkill",
                icon=exit_icon,
                title='Exited!',
                body='Windows usbkill has exited! It will no longer detect changes in the USB ports!' 
            )
    def os_exit():
        os._exit(0)

    exit_toast_thread = threading.Thread(target=exit_toast)
    os_exit_thread = threading.Thread(target=os_exit)
    exit_toast_thread.start()
    time.sleep(1)
    os_exit_thread.start()

# Tray
def create_tray_icon():
    pythoncom.CoInitialize()
    icon = pystray.Icon("Test", image, 'Windows usbkill', menu=pystray.Menu(
        pystray.MenuItem("Pause", toggle_pause, checked=lambda item: is_paused),
        pystray.MenuItem("About", trigger_about_section),
        pystray.MenuItem("Exit", terminate)
))
    icon.run()

# Start
scan_thread = threading.Thread(target=start_scan)
scan_thread.start()
trigger_thread = threading.Thread(target=create_tray_icon)
trigger_thread.start()