Add device model clients and use them

Author: achow101

hwilib/commands.py

@@ -12,14 +12,17 @@
 from .devices import __all__ as all_devs
 
 def get_client_class(device_type):
-    device_type = device_type.split('_')[0]
-    class_name = device_type.capitalize()
-    module = device_type.lower()
+    device_type_split = device_type.split('_')
+    if device_type_split[-1].lower() == 'simulator':
+        del device_type_split[-1]
+    device_type_split = [x.capitalize() for x in device_type_split]
+    device_model = ''.join(device_type_split)
+    module = device_type_split[0].lower()
 
     try:
         imported_dev = importlib.import_module('.devices.' + module, __package__)
-        client_constructor = getattr(imported_dev, class_name + 'Client')
-    except ImportError:
+        client_constructor = getattr(imported_dev, device_model + 'Client')
+    except (ImportError, AttributeError):
         raise UnknownDeviceError('Unknown device type specified')
 
     return client_constructor
@@ -57,7 +60,7 @@ def find_device(password='', device_type=None, fingerprint=None, expert=False):
             continue
         client = None
         try:
-            client = get_client(d['type'], d['path'], password, expert)
+            client = get_client(d['model'], d['path'], password, expert)
 
             master_fpr = d.get('fingerprint', None)
             if master_fpr is None:

hwilib/devices/digitalbitbox.py

@@ -583,6 +583,11 @@ def prompt_pin(self):
     def send_pin(self, pin):
         raise UnavailableActionError('The Digital Bitbox does not need a PIN sent from the host')
 
+class Digitalbitbox01Client(DigitalbitboxClient):
+    def __init__(self, path, password='', expert=False):
+        super(Digitalbitbox01Client, self).__init__(path, password, expert)
+        self.type = 'Digital BitBox01'
+
 def enumerate(password=''):
     results = []
     devices = hid.enumerate(DBB_VENDOR_ID, DBB_DEVICE_ID)

hwilib/devices/ledger.py

@@ -74,6 +74,7 @@ class LedgerClient(HardwareWalletClient):
 
     def __init__(self, path, password='', expert=False):
         super(LedgerClient, self).__init__(path, password, expert)
+        self.type = 'Ledger Nano S and X'
 
         if path.startswith('tcp'):
             split_path = path.split(':')
@@ -354,6 +355,16 @@ def prompt_pin(self):
     def send_pin(self, pin):
         raise UnavailableActionError('The Ledger Nano S and X do not need a PIN sent from the host')
 
+class LedgerNanoSClient(LedgerClient):
+    def __init__(self, path, password='', expert=False):
+        super(LedgerNanoSClient, self).__init__(path, password, expert)
+        self.type = 'Ledger Nano S'
+
+class LedgerNanoXClient(LedgerClient):
+    def __init__(self, path, password='', expert=False):
+        super(LedgerNanoXClient, self).__init__(path, password, expert)
+        self.type = 'Ledger Nano X'
+
 def enumerate(password=''):
     results = []
     devices = []

hwilib/devices/trezor.py

@@ -434,6 +434,16 @@ def send_pin(self, pin):
             return {'success': False}
         return {'success': True}
 
+class Trezor1Client(TrezorClient):
+    def __init__(self, path, password='', expert=False):
+        super(Trezor1Client, self).__init__(path, password, expert)
+        self.type = 'Trezor 1'
+
+class TrezorTClient(TrezorClient):
+    def __init__(self, path, password='', expert=False):
+        super(TrezorTClient, self).__init__(path, password, expert)
+        self.type = 'Trezor T'
+
 def enumerate(password=''):
     results = []
     for dev in enumerate_devices():