usb_collector.py

import winreg

# Common USB Vendor IDs lookup table
USB_VENDOR_IDS = {
    '0000': 'Generic',
    '0781': 'SanDisk',
    '0951': 'Kingston',
    '090C': 'Silicon Motion',
    '13FE': 'Phison Electronics',
    '1F75': 'Innostor Technology',
    '8564': 'Transcend',
    '0930': 'Toshiba',
    '1005': 'Apacer',
    '0D7D': 'Phison Electronics',
    '058F': 'Alcor Micro',
    '1307': 'USBest Technology',
    '048D': 'Integrated Technology Express',
    '054C': 'Sony',
    '04E8': 'Samsung',
    '0424': 'Microchip Technology',
    '05DC': 'Lexar',
    '18A5': 'Verbatim',
    '125F': 'ADATA',
    '1B1C': 'Corsair',
    '0BC2': 'Seagate',
    '1058': 'Western Digital',
    '8087': 'Intel',
    '22B8': 'Motorola',
    '2717': 'Xiaomi',
    '05E3': 'Genesys Logic',
}


def get_usb_history():
    """
    Collect USB storage device history from Windows Registry.
    Sources: USBSTOR for device info, USB for VID/PID and manufacturer
    """
    usb_devices = []
    vid_pid_info = _get_vid_pid_info()
    
    try:
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, 
                            r"SYSTEM\CurrentControlSet\Enum\USBSTOR", 
                            0, winreg.KEY_READ)
        i = 0
        while True:
            try:
                # subkey format: Disk&Ven_VENDOR&Prod_PRODUCT&Rev_REVISION
                subkey_name = winreg.EnumKey(key, i)
                usb_key = winreg.OpenKey(key, subkey_name, 0, winreg.KEY_READ)
                
                # Parse vendor and product from subkey name
                vendor, product = _parse_vendor_product(subkey_name)
                
                j = 0
                while True:
                    try:
                        # This is the serial number
                        serial = winreg.EnumKey(usb_key, j)
                        device_key = winreg.OpenKey(usb_key, serial, 0, winreg.KEY_READ)
                        
                        # Get FriendlyName
                        try:
                            friendly_name, _ = winreg.QueryValueEx(device_key, "FriendlyName")
                        except FileNotFoundError:
                            friendly_name = f"{vendor} {product}".strip() or "Unknown USB Device"
                        
                        # Get VID/PID and manufacturer from USB registry
                        base_serial = serial.split('&')[0] if '&' in serial else serial
                        device_info = vid_pid_info.get(base_serial, {})
                        vid = device_info.get('vid', '')
                        pid = device_info.get('pid', '')
                        usb_manufacturer = device_info.get('manufacturer', '')
                        
                        # Determine best vendor name
                        # Priority: 1) USBSTOR vendor, 2) USB Mfg field, 3) VID lookup, 4) "Unknown"
                        if vendor:
                            final_vendor = vendor
                        elif usb_manufacturer and 'compatible' not in usb_manufacturer.lower():
                            final_vendor = usb_manufacturer
                        elif vid and vid.upper() in USB_VENDOR_IDS:
                            final_vendor = USB_VENDOR_IDS[vid.upper()]
                        else:
                            final_vendor = "Unknown" if vid else ""
                        
                        # Get ContainerID for additional identification
                        try:
                            container_id, _ = winreg.QueryValueEx(device_key, "ContainerID")
                        except FileNotFoundError:
                            container_id = ""
                        
                        usb_devices.append({
                            'name': friendly_name,
                            'vendor': final_vendor,
                            'product': product,
                            'vid': vid,
                            'pid': pid,
                            'serial': serial,
                            'container_id': container_id,
                            'first_insert': '',
                            'last_insert': ''
                        })
                        
                        winreg.CloseKey(device_key)
                        j += 1
                    except OSError:
                        break
                winreg.CloseKey(usb_key)
                i += 1
            except OSError:
                break
        winreg.CloseKey(key)
    except Exception as e:
        print(f"Error reading USBSTOR registry: {e}")
    
    return usb_devices


def _get_vid_pid_info():
    """
    Build a mapping of serial numbers to VID/PID and manufacturer from USB registry.
    USB registry has format: USB\\VID_XXXX&PID_YYYY\\SERIAL
    Returns: {serial: {'vid': '...', 'pid': '...', 'manufacturer': '...'}}
    """
    vid_pid_info = {}
    
    try:
        usb_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                                r"SYSTEM\CurrentControlSet\Enum\USB",
                                0, winreg.KEY_READ)
        i = 0
        while True:
            try:
                # subkey format: VID_XXXX&PID_YYYY
                vid_pid_key_name = winreg.EnumKey(usb_key, i)
                
                # Skip non-VID entries
                if not vid_pid_key_name.startswith("VID_"):
                    i += 1
                    continue
                
                # Parse VID and PID
                vid, pid = _parse_vid_pid(vid_pid_key_name)
                
                vid_pid_subkey = winreg.OpenKey(usb_key, vid_pid_key_name, 0, winreg.KEY_READ)
                j = 0
                while True:
                    try:
                        # This is the serial number
                        serial = winreg.EnumKey(vid_pid_subkey, j)
                        
                        # Try to get manufacturer from device key
                        manufacturer = ""
                        try:
                            device_key = winreg.OpenKey(vid_pid_subkey, serial, 0, winreg.KEY_READ)
                            try:
                                mfg_value, _ = winreg.QueryValueEx(device_key, "Mfg")
                                # Clean up the manufacturer string (remove @file.inf,%key%; prefix)
                                if ';' in mfg_value:
                                    manufacturer = mfg_value.split(';')[-1].strip()
                                else:
                                    manufacturer = mfg_value.strip()
                            except FileNotFoundError:
                                pass
                            winreg.CloseKey(device_key)
                        except OSError:
                            pass
                        
                        vid_pid_info[serial] = {
                            'vid': vid,
                            'pid': pid,
                            'manufacturer': manufacturer
                        }
                        j += 1
                    except OSError:
                        break
                winreg.CloseKey(vid_pid_subkey)
                i += 1
            except OSError:
                break
        winreg.CloseKey(usb_key)
    except Exception as e:
        print(f"Error reading USB registry for VID/PID: {e}")
    
    return vid_pid_info


def _parse_vendor_product(subkey_name):
    """
    Parse vendor and product from USBSTOR subkey name.
    Format: Disk&Ven_VENDOR&Prod_PRODUCT&Rev_REVISION
    """
    vendor = ""
    product = ""
    
    parts = subkey_name.split('&')
    for part in parts:
        if part.startswith('Ven_'):
            vendor = part[4:].strip('_').replace('_', ' ')
        elif part.startswith('Prod_'):
            product = part[5:].strip('_').replace('_', ' ')
    
    return vendor, product


def _parse_vid_pid(vid_pid_key_name):
    """
    Parse VID and PID from USB subkey name.
    Format: VID_XXXX&PID_YYYY
    """
    vid = ""
    pid = ""
    
    parts = vid_pid_key_name.split('&')
    for part in parts:
        if part.startswith('VID_'):
            vid = part[4:]
        elif part.startswith('PID_'):
            pid = part[4:]
    
    return vid, pid


def get_usb_mountpoints():
    """
    Get mounted drive letters for USB devices from MountedDevices registry.
    """
    mountpoints = {}
    
    try:
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                            r"SYSTEM\MountedDevices",
                            0, winreg.KEY_READ)
        i = 0
        while True:
            try:
                name, value, _ = winreg.EnumValue(key, i)
                if name.startswith("\\DosDevices\\") and b"USBSTOR" in value:
                    drive_letter = name.replace("\\DosDevices\\", "")
                    # Extract serial from the mounted device path
                    try:
                        decoded = value.decode('utf-16-le', errors='ignore')
                        if "USBSTOR#Disk" in decoded:
                            # Extract serial number
                            parts = decoded.split('#')
                            if len(parts) >= 3:
                                serial = parts[2].split('&')[0]
                                mountpoints[serial] = drive_letter
                    except:
                        pass
                i += 1
            except OSError:
                break
        winreg.CloseKey(key)
    except Exception as e:
        print(f"Error reading MountedDevices: {e}")
    
    return mountpoints