209 lines
6.0 KiB
Python
209 lines
6.0 KiB
Python
# 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()
|