'''USB HID implementation of Transport.''' import hid import time from .transport import TransportV1, TransportV2, ConnectionError def enumerate(): """ Return a list of available TREZOR devices. """ devices = {} for d in hid.enumerate(0, 0): vendor_id = d['vendor_id'] product_id = d['product_id'] serial_number = d['serial_number'] interface_number = d['interface_number'] path = d['path'] # HIDAPI on Mac cannot detect correct HID interfaces, so device with # DebugLink doesn't work on Mac... if devices.get(serial_number) != None and devices[serial_number][0] == path: raise Exception("Two devices with the same path and S/N found. This is Mac, right? :-/") if (vendor_id, product_id) in DEVICE_IDS: devices.setdefault(serial_number, [None, None]) if interface_number == 0 or interface_number == -1: # normal link devices[serial_number][0] = path elif interface_number == 1: # debug link devices[serial_number][1] = path # List of two-tuples (path_normal, path_debuglink) return list(devices.values()) def path_to_transport(path): try: device = [ d for d in hid.enumerate(0, 0) if d['path'] == path ][0] except IndexError: raise ConnectionError("Connection failed") # VID/PID found, let's find proper transport try: transport = DEVICE_TRANSPORTS[(device['vendor_id'], device['product_id'])] except IndexError: raise Exception("Unknown transport for VID:PID %04x:%04x" % (vid, pid)) return transport class _HidTransport(object): def __init__(self, device, *args, **kwargs): self.hid = None self.hid_version = None device = device[int(bool(kwargs.get('debug_link')))] super(_HidTransport, self).__init__(device, *args, **kwargs) def is_connected(self): """ Check if the device is still connected. """ for d in hid.enumerate(0, 0): if d['path'] == self.device: return True return False def _open(self): self.hid = hid.device() self.hid.open_path(self.device) self.hid.set_nonblocking(True) # determine hid_version r = self.hid.write([0, 63, ] + [0xFF] * 63) if r == 65: self.hid_version = 2 return r = self.hid.write([63, ] + [0xFF] * 63) if r == 64: self.hid_version = 1 return raise ConnectionError("Unknown HID version") def _close(self): self.hid.close() self.hid = None def _write_chunk(self, chunk): if len(chunk) != 64: raise Exception("Unexpected data length") if self.hid_version == 2: self.hid.write(b'\0' + chunk) else: self.hid.write(chunk) def _read_chunk(self): start = time.time() while True: data = self.hid.read(64) if not len(data): if time.time() - start > 10: # Over 10 s of no response, let's check if # device is still alive if not self.is_connected(): raise ConnectionError("Connection failed") # Restart timer start = time.time() time.sleep(0.001) continue break if len(data) != 64: raise Exception("Unexpected chunk size: %d" % len(data)) return bytearray(data) class HidTransportV1(_HidTransport, TransportV1): pass class HidTransportV2(_HidTransport, TransportV2): pass DEVICE_IDS = [ (0x534c, 0x0001), # TREZOR (0x1209, 0x53C0), # TREZORv2 Bootloader (0x1209, 0x53C1), # TREZORv2 ] DEVICE_TRANSPORTS = { (0x534c, 0x0001): HidTransportV1, # TREZOR (0x1209, 0x53C0): HidTransportV2, # TREZORv2 Bootloader (0x1209, 0x53C1): HidTransportV2, # TREZORv2 } # Backward compatible wrapper, decides for proper transport # based on VID/PID of given path def HidTransport(device, *args, **kwargs): transport = path_to_transport(device[0]) return transport(device, *args, **kwargs) # Backward compatibility hack; HidTransport is a function, not a class like before HidTransport.enumerate = enumerate