ws_usbdisable.py

r"""
================================================================================
 ws_usbdisable.py
================================================================================
 Author      : Sparviero
 Created     : 2026-04-14
 Description : Toggles USB mass-storage ports on/off by flipping the USBSTOR
               service Start value in the Windows registry.
               Double-click to run — admin privileges are requested automatically
               via UAC if not already elevated.

 Python commands:
   python -m venv .venv
   .venv\Scripts\activate.bat
   python.exe -m pip install --upgrade pip
   pip install pyinstaller
   pyinstaller --onefile --console ws_usbdisable.py

 Usage (admin required):
   ws_usbdisable.exe   — checks current USB state and toggles it (enable ↔ disable)
                         prompts for an optional immediate system restart

 Registry key : HKLM\SYSTEM\CurrentControlSet\Services\USBSTOR
   Start = 3  — USB mass-storage ENABLED
   Start = 4  — USB mass-storage DISABLED

 Note: affects USB mass-storage devices only (flash drives, external HDDs).
       HID devices (keyboard, mouse) use separate drivers and are not affected.

 Build dependency:
   pyinstaller

 Build (requires Python on the build machine):
   pyinstaller --onefile --console ws_usbdisable.py
   output: dist\ws_usbdisable.exe
================================================================================
"""

import sys
import ctypes
import winreg
import os

USBSTOR_KEY = r"SYSTEM\CurrentControlSet\Services\USBSTOR"
START_VALUE = "Start"
ENABLED  = 3
DISABLED = 4


def is_admin():
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except Exception:
        return False


def relaunch_as_admin():
    script = os.path.abspath(sys.argv[0])
    params = " ".join([f'"{a}"' for a in sys.argv[1:]])
    ctypes.windll.shell32.ShellExecuteW(
        None, "runas", sys.executable, f'"{script}" {params}', None, 1
    )
    sys.exit(0)


def get_usb_state():
    with winreg.OpenKey(
        winreg.HKEY_LOCAL_MACHINE, USBSTOR_KEY, 0, winreg.KEY_READ
    ) as key:
        value, _ = winreg.QueryValueEx(key, START_VALUE)
    return value


def set_usb_state(new_value):
    with winreg.OpenKey(
        winreg.HKEY_LOCAL_MACHINE, USBSTOR_KEY, 0,
        winreg.KEY_SET_VALUE | winreg.KEY_WOW64_64KEY
    ) as key:
        winreg.SetValueEx(key, START_VALUE, 0, winreg.REG_DWORD, new_value)


def main():
    os.system("title USB Port Toggle")

    if not is_admin():
        print("Privilege check failed — relaunching as Administrator...")
        relaunch_as_admin()
        return

    print("=" * 45)
    print("         USB Port Toggle Utility")
    print("=" * 45)

    try:
        current = get_usb_state()
    except PermissionError:
        print("\n[ERROR] Cannot read registry. Run as Administrator.")
        input("\nPress Enter to exit...")
        sys.exit(1)
    except FileNotFoundError:
        print("\n[ERROR] USBSTOR registry key not found.")
        input("\nPress Enter to exit...")
        sys.exit(1)

    if current == ENABLED:
        print(f"\nCurrent status : USB ports are ENABLED  (Start={current})")

        print("Action         : Disabling USB mass-storage...")
        print("\n  Note: only mass-storage devices are affected (flash drives, external HDDs).")
        print("        Keyboards and mice (HID) use separate drivers and will NOT be disabled.")
        set_usb_state(DISABLED)
        print("\n>>> USB mass-storage is now *DISABLED* <<<")
        print("\n[!] A system RESTART is required for the change to take effect.")

    elif current == DISABLED:
        print(f"\nCurrent status : USB ports are DISABLED (Start={current})")
        print("Action         : Enabling USB ports...")
        set_usb_state(ENABLED)
        print("\n>>> USB ports are now *ENABLED* <<<")
        print("\n[!] A system RESTART is required for the change to take effect.")

    else:
        print(f"\n[WARNING] Unexpected Start value: {current}")
        print("The key was not modified. Verify the registry manually.")

    print()
    restart = input("Do you want to restart the system now? [y/N]: ").strip().lower()
    if restart == "y":
        print("Restarting in 10 seconds...")
        os.system("shutdown /r /t 10 /c \"USB toggle — system restart\"")
    else:
        print("Restart skipped. Changes will take effect after the next reboot.")

    input("\nPress Enter to exit...")


if __name__ == "__main__":
    main()