81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
import usb.core
|
|
import usb.util
|
|
from usb import CLASS_MASS_STORAGE
|
|
|
|
from ereuse_devicehub.ereuse_utils.naming import Naming
|
|
|
|
|
|
def plugged_usbs(multiple=True) -> map or dict: # noqa: C901
|
|
"""
|
|
Gets the plugged-in USB Flash drives (pen-drives).
|
|
|
|
If multiple is true, it returns a map, and a dict otherwise.
|
|
|
|
If multiple is false, this method will raise a :class:`.NoUSBFound` if no USB is found.
|
|
"""
|
|
|
|
class FindPenDrives(object):
|
|
# From https://github.com/pyusb/pyusb/blob/master/docs/tutorial.rst
|
|
def __init__(self, class_):
|
|
self._class = class_
|
|
|
|
def __call__(self, device):
|
|
# first, let's check the device
|
|
if device.bDeviceClass == self._class:
|
|
return True
|
|
# ok, transverse all devices to find an
|
|
# interface that matches our class
|
|
for cfg in device:
|
|
# find_descriptor: what's it?
|
|
intf = usb.util.find_descriptor(cfg, bInterfaceClass=self._class)
|
|
# We don't want Card readers
|
|
if intf is not None:
|
|
try:
|
|
product = intf.device.product.lower()
|
|
except ValueError as e:
|
|
if 'langid' in str(e):
|
|
raise OSError(
|
|
'Cannot get "langid". Do you have permissions?'
|
|
)
|
|
else:
|
|
raise e
|
|
if 'crw' not in product and 'reader' not in product:
|
|
return True
|
|
return False
|
|
|
|
def get_pendrive(pen: usb.Device) -> dict:
|
|
if not pen.manufacturer or not pen.product or not pen.serial_number:
|
|
raise UsbDoesNotHaveHid()
|
|
manufacturer = pen.manufacturer.strip() or str(pen.idVendor)
|
|
model = pen.product.strip() or str(pen.idProduct)
|
|
serial_number = pen.serial_number.strip()
|
|
hid = Naming.hid('USBFlashDrive', manufacturer, model, serial_number)
|
|
return {
|
|
'id': hid, # Make live easier to DeviceHubClient by using _id
|
|
'hid': hid,
|
|
'type': 'USBFlashDrive',
|
|
'serialNumber': serial_number,
|
|
'model': model,
|
|
'manufacturer': manufacturer,
|
|
'vendorId': pen.idVendor,
|
|
'productId': pen.idProduct,
|
|
}
|
|
|
|
result = usb.core.find(
|
|
find_all=multiple, custom_match=FindPenDrives(CLASS_MASS_STORAGE)
|
|
)
|
|
if multiple:
|
|
return map(get_pendrive, result)
|
|
else:
|
|
if not result:
|
|
raise NoUSBFound()
|
|
return get_pendrive(result)
|
|
|
|
|
|
class NoUSBFound(Exception):
|
|
pass
|
|
|
|
|
|
class UsbDoesNotHaveHid(Exception):
|
|
pass
|