diff --git a/src/calibre/devices/winusb.py b/src/calibre/devices/winusb.py index ed970369cd..bf7a01573c 100644 --- a/src/calibre/devices/winusb.py +++ b/src/calibre/devices/winusb.py @@ -228,7 +228,9 @@ for line in '''\ code = CR_CODES[code] CR_CODES[name] = code CR_CODE_NAMES[code] = name - +CM_GET_DEVICE_INTERFACE_LIST_PRESENT = 0 +CM_GET_DEVICE_INTERFACE_LIST_ALL_DEVICES = 1 +CM_GET_DEVICE_INTERFACE_LIST_BITS = 1 setupapi = windll.setupapi cfgmgr = windll.CfgMgr32 kernel32 = windll.Kernel32 @@ -285,6 +287,7 @@ CM_Get_Sibling = cwrap('CM_Get_Sibling', CONFIGRET, POINTER(DEVINST), DEVINST, U CM_Get_Device_ID_Size = cwrap('CM_Get_Device_ID_Size', CONFIGRET, POINTER(ULONG), DEVINST, ULONG) CM_Get_Device_ID = cwrap('CM_Get_Device_IDW', CONFIGRET, DEVINST, LPWSTR, ULONG, ULONG) CM_Request_Device_Eject = cwrap('CM_Request_Device_EjectW', CONFIGRET, DEVINST, c_void_p, LPWSTR, ULONG, ULONG, errcheck=config_err_check) +CM_Get_Device_Interface_List = cwrap('CM_Get_Device_Interface_ListW', CONFIGRET, POINTER(GUID), LPCWSTR, LPWSTR, ULONG, ULONG, errcheck=config_err_check) # }}} @@ -505,7 +508,7 @@ def get_drive_letters_for_device(vendor_id, product_id, bcd=None, debug=False): can be either None, in which case it is not tested, or it must be a list or set like object containing bcds. ''' - rbuf = buf = wbuf = None + rbuf = wbuf = None ans = [] # First search for a device matching the specified USB ids @@ -536,42 +539,12 @@ def get_drive_letters_for_device(vendor_id, product_id, bcd=None, debug=False): return ans # Get the device ids for all descendants of the found device - device_ids = set() + sn_map = get_storage_number_map() for devinst in iterdescendants(devinfo.DevInst): devid, wbuf = get_device_id(devinst, buf=wbuf) - device_ids.add(devid.upper().replace(os.sep, '#')) - if debug: - prints('Device ids for vid=0x%x pid=0x%x: %r' % (vendor_id, product_id, device_ids)) - if not device_ids: - return ans - - drive_map = get_all_removable_drives(allow_fixed=True) - if not drive_map: - raise NoRemovableDrives('No removable drives found!') - - # Now look for volumes whose device path contains one of the child device - # ids we found earlier - with get_device_set() as dev_list: - interface_data = SP_DEVICE_INTERFACE_DATA() - interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA) - i = -1 - while True: - i += 1 - if not SetupDiEnumDeviceInterfaces(dev_list, None, byref(GUID_DEVINTERFACE_VOLUME), i, byref(interface_data)): - break - buf, devinfo, devpath = get_device_interface_detail_data(dev_list, byref(interface_data), buf) - devpath = devpath.upper() - matched = False - for q in device_ids: - if q in devpath: - matched = True - break - if debug: - prints('Found volume with device path: %s Matches a device_id: %s' % (devpath, matched)) - if matched: - drive_letter = drive_letter_from_volume_devpath(devpath, drive_map) - if drive_letter: - ans.append(drive_letter) + drive_letter = find_drive(devinst.value, sn_map) + if drive_letter: + ans.append(drive_letter) return ans _devid_pat = None @@ -581,6 +554,48 @@ def devid_pat(): _devid_pat = re.compile(r'VID_([a-f0-9]{4})&PID_([a-f0-9]{4})&REV_([a-f0-9]{4})', re.I) return _devid_pat +def get_storage_number(devpath): + sdn = STORAGE_DEVICE_NUMBER() + handle = CreateFile(devpath, 0, FILE_SHARE_READ | FILE_SHARE_WRITE, None, OPEN_EXISTING, 0, None) + try: + DeviceIoControl(handle, IOCTL_STORAGE_GET_DEVICE_NUMBER, None, 0, byref(sdn), sizeof(STORAGE_DEVICE_NUMBER), None, None) + finally: + CloseHandle(handle) + return sdn.DeviceNumber + +def get_storage_number_map(drive_types=(DRIVE_REMOVABLE, DRIVE_FIXED)): + mask = GetLogicalDrives() + type_map = {letter:GetDriveType(letter + ':' + os.sep) for i, letter in enumerate(string.ascii_uppercase) if mask & (1 << i)} + drives = (letter for letter, dt in type_map.iteritems() if dt in drive_types) + ans = {} + for letter in drives: + try: + ans[get_storage_number('\\\\.\\' + letter + ':')] = letter + except WindowsError: + continue + return ans + + +def find_drive(devinst, storage_number_map): + buf = None + interface_data = SP_DEVICE_INTERFACE_DATA() + interface_data.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA) + with get_device_set(guid=byref(GUID_DEVINTERFACE_DISK)) as dev_list: + i = -1 + while True: + i += 1 + if not SetupDiEnumDeviceInterfaces(dev_list, None, byref(GUID_DEVINTERFACE_DISK), i, byref(interface_data)): + break + buf, devinfo, devpath = get_device_interface_detail_data(dev_list, byref(interface_data), buf=buf) + if devinfo.DevInst == devinst: + try: + storage_number = get_storage_number(devpath) + except WindowsError: + continue + drive_letter = storage_number_map.get(storage_number) + if drive_letter: + return drive_letter + # }}} def get_usb_devices(): # {{{