实现U盘的自动休眠和唤醒功能.

豆虫

[Python] 纯文本查看 复制代码

import tkinter as tk
from tkinter import ttk, messagebox
import threading
import sys
import os

# --- 引入修正后的 PnP 控制代码 (使用 SetupAPI - 最终版 v2) ---
# --- Start of imported and corrected code block ---
import win32api
# import win32con # 不再需要此模块
import ctypes
from ctypes import wintypes

# --- 定义 Windows GUID 结构体 ---
class GUID(ctypes.Structure):
    _fields_ = [
        ("Data1", wintypes.DWORD),
        ("Data2", wintypes.WORD),
        ("Data3", wintypes.WORD),
        ("Data4", wintypes.BYTE * 8),
    ]
# --- 修正 win32con 常量 ---
DICS_DISABLE = 0x00000002
DICS_ENABLE = 0x00000001
# --- /修正 ---

# --- Windows API 定义 ---
# 修正：函数位于 cfgmgr32.dll (仅用于定位)
cfgmgr32 = ctypes.WinDLL('cfgmgr32.dll', use_last_error=True)
# 新增：函数位于 setupapi.dll (用于更改属性)
setupapi = ctypes.WinDLL('setupapi.dll', use_last_error=True)

CM_LOCATE_DEVNODE_NORMAL = 0x00000000
DICS_FLAG_GLOBAL = 0x00000001
DICS_FLAG_CONFIGSPECIFIC = 0x00000002
DIF_PROPERTYCHANGE = 0x00000012

# --- 修正 win32con 常量 ---
DIGCF_ALLCLASSES = 0x00000004
DIGCF_PRESENT = 0x00000002
# --- /修正 ---

CR_SUCCESS = 0x00000000
ERROR_SUCCESS = 0x00000000
ERROR_NO_MORE_ITEMS = 259

# 修正：定义 CONFIGRET 为 ULONG
CONFIGRET = wintypes.ULONG

# 定义必要的结构体和类型
class SP_DEVINFO_DATA(ctypes.Structure):
    _fields_ = [
        ('cbSize', wintypes.DWORD),
        ('ClassGuid', GUID),
        ('DevInst', wintypes.DWORD), # cfgmgr32 使用 DevInst
        ('Reserved', ctypes.c_void_p),
    ]

class SP_CLASSINSTALL_HEADER(ctypes.Structure):
    _fields_ = [
        ('cbSize', wintypes.DWORD),
        ('InstallFunction', wintypes.DWORD),
    ]

class SP_PROPCHANGE_PARAMS(ctypes.Structure):
    _fields_ = [
        ('ClassInstallHeader', SP_CLASSINSTALL_HEADER),
        ('StateChange', wintypes.DWORD),
        ('Scope', wintypes.DWORD),
        ('HwProfile', wintypes.DWORD),
    ]

# --- 定义 SetupAPI 函数原型 ---
def define_setupapi_prototypes():
    """定义 SetupAPI 函数的原型，确保参数和返回类型正确。"""
    global setupapi
    try:
        # SetupDiGetClassDevsW
        setupapi.SetupDiGetClassDevsW.argtypes = [
            ctypes.POINTER(GUID),  # CONST GUID * (None is passed, which is OK)
            wintypes.LPCWSTR,      # PCWSTR
            wintypes.HWND,         # HWND
            wintypes.DWORD,        # DWORD
        ]
        setupapi.SetupDiGetClassDevsW.restype = wintypes.HANDLE

        # SetupDiEnumDeviceInfo
        setupapi.SetupDiEnumDeviceInfo.argtypes = [
            wintypes.HANDLE,           # HDEVINFO
            wintypes.DWORD,            # DWORD
            ctypes.POINTER(SP_DEVINFO_DATA) # PSP_DEVINFO_DATA
        ]
        setupapi.SetupDiEnumDeviceInfo.restype = wintypes.BOOL

        # SetupDiDestroyDeviceInfoList
        setupapi.SetupDiDestroyDeviceInfoList.argtypes = [wintypes.HANDLE]
        setupapi.SetupDiDestroyDeviceInfoList.restype = wintypes.BOOL

        # SetupDiSetClassInstallParamsW - 修改 argtypes 为 c_void_p
        setupapi.SetupDiSetClassInstallParamsW.argtypes = [
            wintypes.HANDLE,                    # HDEVINFO
            ctypes.POINTER(SP_DEVINFO_DATA),   # PSP_DEVINFO_DATA
            ctypes.c_void_p,                    # PSP_CLASSINSTALL_HEADER (as void*)
            wintypes.DWORD                      # DWORD
        ]
        setupapi.SetupDiSetClassInstallParamsW.restype = wintypes.BOOL

        # SetupDiCallClassInstaller
        setupapi.SetupDiCallClassInstaller.argtypes = [
            wintypes.DWORD,                     # DI_FUNCTION
            wintypes.HANDLE,                    # HDEVINFO
            ctypes.POINTER(SP_DEVINFO_DATA)     # PSP_DEVINFO_DATA
        ]
        setupapi.SetupDiCallClassInstaller.restype = wintypes.BOOL

        print("SetupAPI 函数原型定义成功。")
        return True
    except AttributeError as e:
        print(f"SetupAPI 函数原型定义失败: {e}")
        return False

# 立即调用定义函数
if not define_setupapi_prototypes():
    print("严重错误：SetupAPI 函数原型定义失败，程序无法继续。")
    exit(1)
# --- /定义 SetupAPI 函数原型 ---

def get_device_instance_id_from_drive(drive_letter):
    """
    根据驱动器号获取设备实例ID (Instance ID)
    注意：此函数为演示目的，自动查找实例ID非常复杂。
    """
    try:
        import wmi
        c = wmi.WMI()
        for logical_disk in c.Win32_LogicalDisk(DriveType=2, DeviceID=f"{drive_letter}:"):
             for partition in logical_disk.associators(wmi_association_class="Win32_LogicalDiskToPartition"):
                 for disk_drive in partition.associators(wmi_association_class="Win32_DiskDriveToDiskPartition"):
                     print(f"找到关联的磁盘驱动器: {disk_drive.Caption}, PnP ID: {disk_drive.PNPDeviceID}")
                     return disk_drive.PNPDeviceID

        print(f"未能找到驱动器 {drive_letter}: 的关联磁盘驱动器。")
        return None

    except ImportError:
        print("错误: 需要安装 'wmi' 和 'pywin32' 库。请运行 'pip install wmi pywin32'。")
        return None
    except Exception as e:
        print(f"获取实例ID时出错: {e}")
        return None

def locate_device_by_instance_id(instance_id):
    """使用 cfgmgr32 定位设备，返回 DevInst。"""
    cm_locate_devnode = cfgmgr32.CM_Locate_DevNodeW
    cm_locate_devnode.argtypes = [ctypes.POINTER(wintypes.DWORD), wintypes.LPCWSTR, wintypes.ULONG]
    cm_locate_devnode.restype = CONFIGRET

    dev_inst = wintypes.DWORD()
    res = cm_locate_devnode(ctypes.byref(dev_inst), instance_id, CM_LOCATE_DEVNODE_NORMAL)
    if res != CR_SUCCESS:
        print(f"定位设备 '{instance_id}' 失败。错误代码: {res}")
        return None
    print(f"成功定位设备 '{instance_id}'，DevInst: {dev_inst.value}")
    return dev_inst.value

def get_device_class_guid(dev_inst):
    """使用 SetupAPI 枚举设备来获取 ClassGuid。"""
    hdevinfo = setupapi.SetupDiGetClassDevsW(
        None,                           # lpClassGuid (None for all classes)
        None,                           # Enumerator (None for all)
        0,                              # hwndParent
        DIGCF_ALLCLASSES | DIGCF_PRESENT # Flags
    )

    if hdevinfo == ctypes.c_void_p(-1).value: # INVALID_HANDLE_VALUE
        last_error = ctypes.GetLastError()
        print(f"SetupDiGetClassDevsW 失败。错误代码: {last_error}")
        return None
    print(f"SetupDiGetClassDevsW 成功，句柄: {hdevinfo}")

    dev_info_data = SP_DEVINFO_DATA()
    dev_info_data.cbSize = ctypes.sizeof(SP_DEVINFO_DATA)
    index = 0

    while True:
        if not setupapi.SetupDiEnumDeviceInfo(hdevinfo, index, ctypes.byref(dev_info_data)):
            last_error = ctypes.GetLastError()
            if last_error == ERROR_NO_MORE_ITEMS:
                print("SetupDiEnumDeviceInfo: 列表末尾。")
                break # No more devices
            else:
                print(f"SetupDiEnumDeviceInfo 错误: {last_error}")
                setupapi.SetupDiDestroyDeviceInfoList(hdevinfo)
                return None # Error occurred

        # 比较 DevInst
        current_devinst = dev_info_data.DevInst
        print(f"  检查索引 {index}，DevInst: {current_devinst}")
        if current_devinst == dev_inst:
            class_guid = dev_info_data.ClassGuid
            setupapi.SetupDiDestroyDeviceInfoList(hdevinfo)
            guid_str = f"{{{class_guid.Data1:08X}-{class_guid.Data2:04X}-{class_guid.Data3:04X}-{''.join('%02X' % b for b in class_guid.Data4[:2])}-{''.join('%02X' % b for b in class_guid.Data4[2:])}}}"
            print(f"找到设备 {dev_inst} 的 ClassGuid: {guid_str}")
            return class_guid

        index += 1

    setupapi.SetupDiDestroyDeviceInfoList(hdevinfo)
    print(f"未能找到设备实例 {dev_inst} 的 ClassGuid。")
    return None

def change_device_state_setupapi(dev_inst, enable=True):
    """使用 SetupAPI 更改设备状态。需要管理员权限。"""
    print(f"[SetupAPI] 开始 {'启用' if enable else '禁用'} 设备，DevInst: {dev_inst}")
    # 1. 获取设备的 ClassGuid
    class_guid = get_device_class_guid(dev_inst)
    if class_guid is None:
        print("[SetupAPI] 无法获取设备的 ClassGuid，SetupAPI 操作失败。")
        return False

    # 2. 创建一个针对该 ClassGuid 的设备信息集
    hdevinfo = setupapi.SetupDiGetClassDevsW(
        ctypes.byref(class_guid),       # lpClassGuid
        None,                           # Enumerator (None for all matching class)
        0,                              # hwndParent
        DIGCF_PRESENT                   # Only present devices in this class
    )

    if hdevinfo == ctypes.c_void_p(-1).value: # INVALID_HANDLE_VALUE
        last_error = ctypes.GetLastError()
        print(f"[SetupAPI] SetupDiGetClassDevsW (按类获取) 失败。错误代码: {last_error}")
        return False
    print(f"[SetupAPI] SetupDiGetClassDevsW (按类) 成功，句柄: {hdevinfo}")

    # 3. 在这个列表中找到 DevInst 匹配的设备
    dev_info_data = SP_DEVINFO_DATA()
    dev_info_data.cbSize = ctypes.sizeof(SP_DEVINFO_DATA)
    index = 0

    found = False
    while True:
        if not setupapi.SetupDiEnumDeviceInfo(hdevinfo, index, ctypes.byref(dev_info_data)):
            last_error = ctypes.GetLastError()
            if last_error == ERROR_NO_MORE_ITEMS:
                print(f"[SetupAPI] SetupDiEnumDeviceInfo (按类查找): 列表末尾。")
                break
            else:
                print(f"[SetupAPI] SetupDiEnumDeviceInfo (按类查找) 错误: {last_error}")
                setupapi.SetupDiDestroyDeviceInfoList(hdevinfo)
                return False
        if dev_info_data.DevInst == dev_inst:
            found = True
            print(f"[SetupAPI] 在类别列表中找到目标设备，索引: {index}")
            break
        index += 1

    if not found:
        print(f"[SetupAPI] 在 ClassGuid {class_guid} 的设备列表中找不到 DevInst {dev_inst}。")
        setupapi.SetupDiDestroyDeviceInfoList(hdevinfo)
        return False

    # 4. 设置属性更改参数
    prop_change_params = SP_PROPCHANGE_PARAMS()
    prop_change_params.ClassInstallHeader.cbSize = ctypes.sizeof(SP_CLASSINSTALL_HEADER)
    prop_change_params.ClassInstallHeader.InstallFunction = DIF_PROPERTYCHANGE
    # --- 修正 win32con 常量 ---
    prop_change_params.StateChange = DICS_DISABLE if not enable else DICS_ENABLE
    # --- /修正 ---
    prop_change_params.Scope = DICS_FLAG_GLOBAL
    prop_change_params.HwProfile = 0 # 当前硬件配置文件

    # 5. 设置类安装参数 - 调用时不变，但 argtypes 已修改
    if not setupapi.SetupDiSetClassInstallParamsW(
        hdevinfo,
        ctypes.byref(dev_info_data),
        ctypes.byref(prop_change_params), # 传递整个结构体的指针
        ctypes.sizeof(SP_PROPCHANGE_PARAMS)
    ):
        last_error = ctypes.GetLastError()
        print(f"[SetupAPI] 设置类安装参数失败。错误代码: {last_error}")
        setupapi.SetupDiDestroyDeviceInfoList(hdevinfo)
        return False
    print(f"[SetupAPI] 类安装参数设置成功。")

    # 6. 调用更改
    if not setupapi.SetupDiCallClassInstaller(
        DIF_PROPERTYCHANGE,
        hdevinfo,
        ctypes.byref(dev_info_data)
    ):
        last_error = ctypes.GetLastError()
        print(f"[SetupAPI] 调用类安装程序失败。错误代码: {last_error}")
        setupapi.SetupDiDestroyDeviceInfoList(hdevinfo)
        return False
    print(f"[SetupAPI] 类安装程序调用成功。")

    setupapi.SetupDiDestroyDeviceInfoList(hdevinfo)
    action = "启用" if enable else "禁用"
    print(f"[SetupAPI] 成功{action}设备。")
    return True


def disable_device(instance_id):
    """使用实例ID禁用设备。需要管理员权限。"""
    print(f"[Main] 开始禁用设备: {instance_id}")
    dev_inst = locate_device_by_instance_id(instance_id)
    if dev_inst is None:
        return False

    # 尝试使用 cfgmgr32 (可能在某些系统上可用)
    try:
        cm_disable_devinst = cfgmgr32.CM_Disable_DevInst
        cm_disable_devinst.argtypes = [wintypes.DWORD, wintypes.ULONG]
        cm_disable_devinst.restype = CONFIGRET

        res = cm_disable_devinst(dev_inst, 0) # 0 for default flags
        if res == CR_SUCCESS:
            print(f"[Main] 成功禁用设备: {instance_id} (通过 CM_Disable_DevInst)")
            return True
        else:
            print(f"[Main] CM_Disable_DevInst 失败，错误代码: {res}。尝试 SetupAPI 方法。")
    except AttributeError:
        print("[Main] CM_Disable_DevInst 函数不可用。尝试 SetupAPI 方法。")

    # 如果 cfgmgr32 失败或不可用，使用 SetupAPI
    return change_device_state_setupapi(dev_inst, enable=False)

def enable_device(instance_id):
    """使用实例ID启用设备。需要管理员权限。"""
    print(f"[Main] 开始启用设备: {instance_id}")
    dev_inst = locate_device_by_instance_id(instance_id)
    if dev_inst is None:
        return False

    # 尝试使用 cfgmgr32 (可能在某些系统上可用)
    try:
        cm_enable_devinst = cfgmgr32.CM_Enable_DevInst
        cm_enable_devinst.argtypes = [wintypes.DWORD, wintypes.ULONG]
        cm_enable_devinst.restype = CONFIGRET

        res = cm_enable_devinst(dev_inst, 0) # 0 for default flags
        if res == CR_SUCCESS:
            print(f"[Main] 成功启用设备: {instance_id} (通过 CM_Enable_DevInst)")
            return True
        else:
            print(f"[Main] CM_Enable_DevInst 失败，错误代码: {res}。尝试 SetupAPI 方法。")
    except AttributeError:
        print("[Main] CM_Enable_DevInst 函数不可用。尝试 SetupAPI 方法。")

    # 如果 cfgmgr32 失败或不可用，使用 SetupAPI
    return change_device_state_setupapi(dev_inst, enable=True)

# --- End of imported and corrected code block ---

class USBControllerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("U盘控制器 (PnP - SetupAPI 终版 v2)")
        self.root.geometry("500x250") # 增加宽度以容纳长ID

        # 提示标签
        info_label = tk.Label(root, text="请确保以管理员身份运行此程序！", fg="red")
        info_label.pack(pady=(10, 5))

        # U盘盘符输入框
        tk.Label(root, text="U盘盘符 (例如 F):").pack(pady=(5, 0))
        self.drive_entry = tk.Entry(root, width=10)
        self.drive_entry.insert(0, "F")  # 默认值为您提供的F盘
        self.drive_entry.pack()

        # 实例ID输入框
        tk.Label(root, text="设备实例ID (建议直接粘贴):").pack(pady=(10, 0))
        self.id_entry = tk.Entry(root, width=70) # 增加宽度
        # 将您查到的ID预填入输入框，注意格式与设备管理器中的一致
        # 设备管理器: USBSTOR\Disk&Ven_VendorCo&Prod_ProductCode&Rev_2.00\1664551093628085247&0
        # Python字符串中需要转义反斜杠: USBSTOR\\Disk&Ven_VendorCo&Prod_ProductCode&Rev_2.00\\1664551093628085247&0
        # 也尝试使用 WMI 获取到的大写格式
        self.id_entry.insert(0, "USBSTOR\\Disk&Ven_VendorCo&Prod_ProductCode&Rev_2.00\\1664551093628085247&0")
        self.id_entry.pack(pady=(0, 10))

        # 按钮框架
        button_frame = tk.Frame(root)
        button_frame.pack(pady=10)

        # 打开 (唤醒/启用) 按钮
        self.open_button = tk.Button(button_frame, text="唤醒 (启用)", command=self.wake_up_usb, bg="#90EE90", width=12)
        self.open_button.pack(side=tk.LEFT, padx=5)

        # 关闭 (休眠/禁用) 按钮
        self.close_button = tk.Button(button_frame, text="休眠 (禁用)", command=self.sleep_usb, bg="#FFB6C1", width=12)
        self.close_button.pack(side=tk.LEFT, padx=5)

        # 状态标签
        self.status_label = tk.Label(root, text="就绪", fg="blue")
        self.status_label.pack(pady=(10, 0))

        # 添加一个按钮用于触发 get_device_instance_id_from_drive 测试
        test_button = tk.Button(root, text="测试获取实例ID (F盘)", command=lambda: self.test_get_id('F'))
        test_button.pack(pady=(5, 0))

    def test_get_id(self, drive):
        """测试获取实例ID的函数"""
        print(f"--- 开始测试 get_device_instance_id_from_drive('{drive}') ---")
        id_found = get_device_instance_id_from_drive(drive)
        if id_found:
            print(f"成功获取到ID: {id_found}")
            # 可选：自动填充到输入框
            self.id_entry.delete(0, tk.END)
            self.id_entry.insert(0, id_found)
            self.status_label.config(text="测试获取ID成功", fg="green")
        else:
            print("获取ID失败或未找到。")
            self.status_label.config(text="测试获取ID失败", fg="red")
        print("--- 测试结束 ---")


    def _run_in_thread(self, func, *args):
        """在后台线程中运行函数，避免阻塞GUI"""
        def target():
            try:
                result = func(*args)
                self.root.after(0, lambda: self._update_status(result, func.__name__))
            except Exception as e:
                error_msg = str(e)
                self.root.after(0, lambda msg=error_msg: self._show_error(msg))
        
        thread = threading.Thread(target=target, daemon=True)
        thread.start()

    def _update_status(self, success, operation_name):
        """更新GUI状态"""
        op_map = {"wake_up_usb": "唤醒", "sleep_usb": "休眠"}
        action = op_map.get(operation_name.replace('_usb', ''), operation_name)
        if success:
            self.status_label.config(text=f"U盘 {action} 操作成功！", fg="green")
        else:
            self.status_label.config(text=f"U盘 {action} 操作失败！", fg="red")

    def _show_error(self, message):
        """显示错误信息"""
        self.status_label.config(text="操作失败", fg="red")
        messagebox.showerror("错误", message)

    def _get_instance_id(self):
        """获取实例ID，优先从输入框，其次尝试自动获取"""
        instance_id = self.id_entry.get().strip()
        if not instance_id:
            drive_letter = self.drive_entry.get().strip()
            if len(drive_letter) != 1:
                messagebox.showwarning("警告", "请输入有效的单个盘符字母。")
                return None
            print(f"正在尝试自动查找驱动器 {drive_letter}: 的实例ID...")
            self.status_label.config(text="正在自动查找实例ID...", fg="orange")
            instance_id = get_device_instance_id_from_drive(drive_letter)
            if instance_id:
                # 可选：将自动找到的ID填充到输入框
                self.id_entry.delete(0, tk.END)
                self.id_entry.insert(0, instance_id)
                print(f"自动找到ID: {instance_id}")
            else:
                messagebox.showerror("错误", f"无法自动找到驱动器 {drive_letter}: 的实例ID。请手动输入。")
                self.status_label.config(text="查找实例ID失败", fg="red")
                return None
        return instance_id

    def sleep_usb(self):
        """休眠 (禁用) U盘"""
        instance_id = self._get_instance_id()
        if not instance_id:
            return
        
        self.status_label.config(text="正在休眠 U盘...", fg="orange")
        self._run_in_thread(lambda: disable_device(instance_id))

    def wake_up_usb(self):
        """唤醒 (启用) U盘"""
        instance_id = self._get_instance_id()
        if not instance_id:
            return

        self.status_label.config(text="正在唤醒 U盘...", fg="orange")
        self._run_in_thread(lambda: enable_device(instance_id))

def check_admin():
    """检查是否以管理员身份运行"""
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False

if __name__ == "__main__":
    if not check_admin():
        # 如果没有管理员权限，则尝试重新启动并请求权限
        print("此程序需要管理员权限才能控制USB设备。正在尝试以管理员身份重启...")
        # 重新构造命令行参数
        params = ' '.join([f'"{arg}"' for arg in sys.argv])
        # 使用 ShellExecuteEx 请求提升权限
        ctypes.windll.shell32.ShellExecuteW(
            None, "runas", sys.executable, f'"{os.path.abspath(__file__)}"', None, 1
        )
        sys.exit(0) # 退出当前非管理员进程

    # 以管理员权限运行时的主流程
    print("已获得管理员权限，启动 GUI...")
    root = tk.Tk()
    app = USBControllerApp(root)
    root.mainloop()

通过 SetupAPI 成功获取设备的类别 GUID {4D36E967-E325-11CE-BFC1-08002BE10318} (磁盘驱动器)。
使用 SetupAPI 的 PropertyChange 功能来禁用/启用设备，这是 Windows 推荐的标准方法。
写在最后的话:
AI真的是太强大了，经过了很多次失败的尝试，最后找到了完美的方案，这个项目展示了底层 Windows API 操作的复杂性，特别是当涉及到 ctypes 与 C 结构体交互时。
如果是人来写的话，要翻多少资料才能完成.

代码经测试，可以完美运行,系统WIN10。


--------------------已打包成exe文件，下面是地址-------------------------


通过网盘分享的文件：u盘控制器.rar
链接: https://pan.baidu.com/s/1MXUxsYwqfGsVWUxPqRSpLw?pwd=5f2g 提取码: 5f2g
--------------------------------------------------------------------------------

豆虫

代码实现的初衷是想U盘长时间插在电脑上，但又怕对U盘寿命有影响，在要往里存资料时可以唤醒,不需要通过拔插U盘来实现。

123tiankongli

U盘的自动休眠和唤醒功能，很牛逼感谢分享

xiaomt

好用吗 有人测试吗

afti

这个有意思

非法操作

很有想法的创作