Skip to content

Commit ce2fd26

Browse files
achow101achow101
committed
Move device enumeration to device specific code
Instead of enumerating only HID devices globally, have each device specific section do its own device enumeration so that we can get devices which use other protocols without effecting global device enumeration code.
1 parent c625e03 commit ce2fd26

File tree

7 files changed

+157
-55
lines changed

7 files changed

+157
-55
lines changed

‎hwilib/commands.py‎

Lines changed: 6 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -53,41 +53,12 @@ def get_client(device_type, device_path, password=None):
5353
# Get a list of all available hardware wallets
5454
def enumerate(args):
5555
result = []
56-
devices = hid.enumerate()
57-
for d in devices:
58-
d_data = {}
59-
# Get trezors
60-
if (d['vendor_id'], d['product_id']) in trezor_device_ids:
61-
d_data['type'] = 'trezor'
62-
# Get keepkeys
63-
elif (d['vendor_id'], d['product_id']) in keepkey_device_ids:
64-
d_data['type'] = 'keepkey'
65-
# Get ledgers, only take interface 0 to avoid HID keyboard
66-
elif (d['vendor_id'], d['product_id']) in ledger_device_ids and \
67-
('interface_number' in d and d['interface_number'] == 0 \
68-
or ('usage_page' in d and d['usage_page'] == 0xffa0)):
69-
d_data['type'] = 'ledger'
70-
# Get DigitalBitboxes
71-
elif (d['vendor_id'], d['product_id']) in digitalbitbox_device_ids:
72-
d_data['type'] = 'digitalbitbox'
73-
# Get ColdCards
74-
elif (d['vendor_id'], d['product_id']) in coldcard_device_ids:
75-
d_data['type'] = 'coldcard'
76-
else:
77-
continue
78-
d_data['path'] = d['path'].decode("utf-8")
79-
d_data['serial_number'] = d['serial_number']
80-
81-
try:
82-
client = get_client(d_data['type'], d_data['path'], args.password)
83-
master_xpub = client.get_pubkey_at_path('m/0h')['xpub']
84-
d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub)
85-
client.close()
86-
except Exception as e:
87-
d_data['error'] = "Could not open client or get fingerprint information: " + str(e)
88-
pass
89-
90-
result.append(d_data)
56+
from .devices.trezor import enumerate as enumerate_trezor
57+
from .devices.coldcard import enumerate as enumerate_coldcard
58+
from .devices.digitalbitbox import enumerate as enumerate_digitalbitbox
59+
from .devices.keepkey import enumerate as enumerate_keepkey
60+
from .devices.ledger import enumerate as enumerate_ledger
61+
result = enumerate_trezor(args.password) + enumerate_coldcard(args.password) + enumerate_digitalbitbox(args.password) + enumerate_keepkey(args.password) + enumerate_ledger(args.password)
9162
return result
9263

9364
# Fingerprint or device type required

‎hwilib/devices/coldcard.py‎

Lines changed: 48 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
# Trezor interaction script
22

33
from ..hwwclient import HardwareWalletClient
4-
from ckcc.client import ColdcardDevice
4+
from ckcc.client import ColdcardDevice, COINKITE_VID, CKCC_PID
55
from ckcc.protocol import CCProtocolPacker
66
from ckcc.constants import MAX_BLK_LEN
7-
from ..base58 import xpub_main_2_test
7+
from ..base58 import xpub_main_2_test, get_xpub_fingerprint_hex
88
from hashlib import sha256
99

1010
import base64
@@ -13,14 +13,21 @@
1313
import io
1414
import time
1515

16+
CC_SIMULATOR_SOCK = '/tmp/ckcc-simulator.sock'
17+
# Using the simulator: https://github.com/Coldcard/firmware/blob/master/unix/README.md
18+
1619
# This class extends the HardwareWalletClient for ColdCard specific things
1720
class ColdCardClient(HardwareWalletClient):
1821

1922
def __init__(self, path, password=''):
2023
super(ColdCardClient, self).__init__(path, password)
21-
device = hid.device()
22-
device.open_path(path.encode())
23-
self.device = ColdcardDevice(dev=device)
24+
# Simulator hard coded pipe socket
25+
if path == CC_SIMULATOR_SOCK:
26+
self.device = ColdcardDevice(sn=path)
27+
else:
28+
device = hid.device()
29+
device.open_path(path.encode())
30+
self.device = ColdcardDevice(dev=device)
2431

2532
# Must return a dict with the xpub
2633
# Retrieves the public key at the specified BIP 32 derivation path
@@ -108,3 +115,39 @@ def wipe_device(self):
108115
# Close the device
109116
def close(self):
110117
self.device.close()
118+
119+
def enumerate(password=None):
120+
results = []
121+
for d in hid.enumerate(COINKITE_VID, CKCC_PID):
122+
d_data = {}
123+
124+
path = d['path'].decode()
125+
d_data['type'] = 'coldcard'
126+
d_data['path'] = path
127+
128+
try:
129+
client = ColdCardClient(path)
130+
master_xpub = client.get_pubkey_at_path('m/0h')['xpub']
131+
d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub)
132+
client.close()
133+
except Exception as e:
134+
d_data['error'] = "Could not open client or get fingerprint information: " + str(e)
135+
136+
results.append(d_data)
137+
# Check if the simulator is there
138+
try:
139+
client = ColdCardClient(CC_SIMULATOR_SOCK)
140+
master_xpub = client.get_pubkey_at_path('m/0h')['xpub']
141+
142+
d_data = {}
143+
d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub)
144+
d_data['type'] = 'coldcard'
145+
d_data['path'] = CC_SIMULATOR_SOCK
146+
results.append(d_data)
147+
client.close()
148+
except RuntimeError as e:
149+
if str(e) == 'Cannot connect to simulator. Is it running?':
150+
pass
151+
else:
152+
raise e
153+
return results

‎hwilib/devices/digitalbitbox.py‎

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@
1010
import binascii
1111
import logging
1212

13-
from ..hwwclient import HardwareWalletClient
13+
from ..hwwclient import HardwareWalletClient, NoPasswordError
1414
from ..serializations import CTransaction, PSBT, hash256, hash160, ser_sig_der, ser_sig_compact, ser_compact_size
15-
from ..base58 import get_xpub_fingerprint, decode, to_address, xpub_main_2_test
15+
from ..base58 import get_xpub_fingerprint, decode, to_address, xpub_main_2_test, get_xpub_fingerprint_hex
1616

1717
applen = 225280 # flash size minus bootloader length
1818
chunksize = 8*512
@@ -23,6 +23,8 @@
2323
HWW_CID = 0xFF000000
2424
HWW_CMD = 0x80 + 0x40 + 0x01
2525

26+
DBB_VENDOR_ID = 0x03eb
27+
DBB_DEVICE_ID = 0x2402
2628

2729
def aes_encrypt_with_iv(key, iv, data):
2830
aes_cbc = pyaes.AESModeOfOperationCBC(key, iv=iv)
@@ -139,10 +141,6 @@ def send_encrypt(msg, password, device):
139141
reply = {'error':'Exception caught while sending encrypted message to DigitalBitbox ' + str(e)}
140142
return reply
141143

142-
class NoPasswordError(Exception):
143-
def __init__(self,*args,**kwargs):
144-
Exception.__init__(self,*args,**kwargs)
145-
146144
# This class extends the HardwareWalletClient for Digital Bitbox specific things
147145
class DigitalBitboxClient(HardwareWalletClient):
148146

@@ -370,3 +368,25 @@ def wipe_device(self):
370368
# Close the device
371369
def close(self):
372370
self.device.close()
371+
372+
def enumerate(password=None):
373+
results = []
374+
for d in hid.enumerate(DBB_VENDOR_ID, DBB_DEVICE_ID):
375+
if ('interface_number' in d and d['interface_number'] == 0 \
376+
or ('usage_page' in d and d['usage_page'] == 0xffff)):
377+
d_data = {}
378+
379+
path = d['path'].decode()
380+
d_data['type'] = 'digitalbitbox'
381+
d_data['path'] = path
382+
383+
try:
384+
client = DigitalBitboxClient(path, password)
385+
master_xpub = client.get_pubkey_at_path('m/0h')['xpub']
386+
d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub)
387+
client.close()
388+
except Exception as e:
389+
d_data['error'] = "Could not open client or get fingerprint information: " + str(e)
390+
391+
results.append(d_data)
392+
return results

‎hwilib/devices/keepkey.py‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,15 @@
66
from keepkeylib import tools
77
from keepkeylib import messages_pb2, types_pb2 as proto
88
from keepkeylib.tx_api import TxApi
9-
from ..base58 import get_xpub_fingerprint, decode, to_address, xpub_main_2_test
9+
from ..base58 import get_xpub_fingerprint, decode, to_address, xpub_main_2_test, get_xpub_fingerprint_hex
1010
from ..serializations import ser_uint256, uint256_from_str
1111

1212
import binascii
1313
import json
1414

15+
KEEPKEY_VENDOR_ID = 0x2B24
16+
KEEPKEY_DEVICE_ID = 0x0001
17+
1518
class TxAPIPSBT(TxApi):
1619

1720
def __init__(self, psbt):
@@ -50,12 +53,8 @@ class KeepKeyClient(HardwareWalletClient):
5053
def __init__(self, path, password=''):
5154
super(KeepKeyClient, self).__init__(path, password)
5255
devices = HidTransport.enumerate()
53-
self.client = None
54-
for d in devices:
55-
if d[0].decode() == path:
56-
transport = HidTransport(d)
57-
self.client = KeepKey(transport)
58-
break
56+
transport = HidTransport((path.encode(), None))
57+
self.client = KeepKey(transport)
5958

6059
# if it wasn't able to find a client, throw an error
6160
if not self.client:
@@ -195,3 +194,23 @@ def wipe_device(self):
195194
# Close the device
196195
def close(self):
197196
self.client.close()
197+
198+
def enumerate(password=None):
199+
results = []
200+
for d in HidTransport.enumerate():
201+
d_data = {}
202+
203+
path = d[0].decode()
204+
d_data['type'] = 'keepkey'
205+
d_data['path'] = path
206+
207+
try:
208+
client = KeepKeyClient(path, password)
209+
master_xpub = client.get_pubkey_at_path('m/0h')['xpub']
210+
d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub)
211+
client.close()
212+
except Exception as e:
213+
d_data['error'] = "Could not open client or get fingerprint information: " + str(e)
214+
215+
results.append(d_data)
216+
return results

‎hwilib/devices/ledger.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,14 @@
77
import json
88
import struct
99
from .. import base58
10+
from ..base58 import get_xpub_fingerprint_hex
1011
from ..serializations import hash256, hash160, ser_uint256, PSBT, CTransaction, HexToBase64
1112
import binascii
1213
import logging
1314

15+
LEDGER_VENDOR_ID = 0x2c97
16+
LEDGER_DEVICE_ID = 0x0001
17+
1418
# This class extends the HardwareWalletClient for Ledger Nano S specific things
1519
class LedgerClient(HardwareWalletClient):
1620

@@ -254,3 +258,25 @@ def wipe_device(self):
254258
# Close the device
255259
def close(self):
256260
self.dongle.close()
261+
262+
def enumerate(password=None):
263+
results = []
264+
for d in hid.enumerate(LEDGER_VENDOR_ID, LEDGER_DEVICE_ID):
265+
if ('interface_number' in d and d['interface_number'] == 0 \
266+
or ('usage_page' in d and d['usage_page'] == 0xffa0)):
267+
d_data = {}
268+
269+
path = d['path'].decode()
270+
d_data['type'] = 'ledger'
271+
d_data['path'] = path
272+
273+
try:
274+
client = LedgerClient(path, password)
275+
master_xpub = client.get_pubkey_at_path('m/0h')['xpub']
276+
d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub)
277+
client.close()
278+
except Exception as e:
279+
d_data['error'] = "Could not open client or get fingerprint information: " + str(e)
280+
281+
results.append(d_data)
282+
return results

‎hwilib/devices/trezor.py‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22

33
from ..hwwclient import HardwareWalletClient
44
from trezorlib.client import TrezorClient as Trezor
5-
from trezorlib.transport import get_transport
5+
from trezorlib.transport import enumerate_devices, get_transport
66
from trezorlib import protobuf, tools
77
from trezorlib import messages as proto
88
from trezorlib.tx_api import TxApi
9-
from ..base58 import get_xpub_fingerprint, decode, to_address, xpub_main_2_test
9+
from ..base58 import get_xpub_fingerprint, decode, to_address, xpub_main_2_test, get_xpub_fingerprint_hex
1010
from ..serializations import ser_uint256, uint256_from_str
1111
from .. import bech32
1212

@@ -193,3 +193,22 @@ def wipe_device(self):
193193
# Close the device
194194
def close(self):
195195
self.client.close()
196+
197+
def enumerate(password=None):
198+
results = []
199+
for dev in enumerate_devices():
200+
d_data = {}
201+
202+
d_data['type'] = 'trezor'
203+
d_data['path'] = dev.get_path()
204+
205+
try:
206+
client = TrezorClient(d_data['path'], password)
207+
master_xpub = client.get_pubkey_at_path('m/0h')['xpub']
208+
d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub)
209+
client.close()
210+
except Exception as e:
211+
d_data['error'] = "Could not open client or get fingerprint information: " + str(e)
212+
213+
results.append(d_data)
214+
return results

‎hwilib/hwwclient.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,7 @@ def wipe_device(self):
4545
def close(self):
4646
raise NotImplementedError('The HardwareWalletClient base class does not '
4747
'implement this method')
48+
49+
class NoPasswordError(Exception):
50+
def __init__(self,*args,**kwargs):
51+
Exception.__init__(self,*args,**kwargs)

0 commit comments

Comments
 (0)