diff --git a/hwilib/devices/coldcard.py b/hwilib/devices/coldcard.py index 8d4cbd8..c0c8735 100644 --- a/hwilib/devices/coldcard.py +++ b/hwilib/devices/coldcard.py @@ -183,6 +183,14 @@ class ColdcardClient(HardwareWalletClient): def close(self): self.device.close() + # Prompt pin + def prompt_pin(self): + raise UnavailableActionError('The Coldcard does not need a PIN sent from the host') + + # Send pin + def send_pin(self): + raise UnavailableActionError('The Coldcard does not need a PIN sent from the host') + def enumerate(password=''): results = [] for d in hid.enumerate(COINKITE_VID, CKCC_PID): diff --git a/hwilib/devices/digitalbitbox.py b/hwilib/devices/digitalbitbox.py index 8c73a2c..bf120ec 100644 --- a/hwilib/devices/digitalbitbox.py +++ b/hwilib/devices/digitalbitbox.py @@ -469,6 +469,14 @@ class DigitalbitboxClient(HardwareWalletClient): def close(self): self.device.close() + # Prompt pin + def prompt_pin(self): + raise UnavailableActionError('The Digtal Bitbox does not need a PIN sent from the host') + + # Send pin + def send_pin(self): + raise UnavailableActionError('The Digital Bitbox does not need a PIN sent from the host') + def enumerate(password=''): results = [] devices = hid.enumerate(DBB_VENDOR_ID, DBB_DEVICE_ID) diff --git a/hwilib/devices/keepkey.py b/hwilib/devices/keepkey.py index 425fafd..1b8ae0a 100644 --- a/hwilib/devices/keepkey.py +++ b/hwilib/devices/keepkey.py @@ -1,12 +1,11 @@ # KeepKey interaction script -from ..hwwclient import HardwareWalletClient, UnavailableActionError +from ..hwwclient import DeviceAlreadyUnlockedError, HardwareWalletClient, UnavailableActionError, DeviceNotReadyError from keepkeylib.transport_hid import HidTransport from keepkeylib.transport_udp import UDPTransport -from keepkeylib.client import KeepKeyClient as KeepKey -from keepkeylib.client import KeepKeyDebugClient as KeepKeyDebug +from keepkeylib.client import BaseClient, DebugWireMixin, DebugLinkMixin, ProtocolMixin, TextUIMixin from keepkeylib import tools -from keepkeylib import messages_pb2, types_pb2 as proto +from keepkeylib import messages_pb2 as messages, types_pb2 as proto from keepkeylib.tx_api import TxApi from ..base58 import get_xpub_fingerprint, decode, to_address, xpub_main_2_test, get_xpub_fingerprint_hex from ..serializations import ser_uint256, uint256_from_str @@ -16,12 +15,20 @@ import base64 import binascii import json import os +import sys KEEPKEY_VENDOR_ID = 0x2B24 KEEPKEY_DEVICE_ID = 0x0001 py_enumerate = enumerate # Need to use the enumerate built-in but there's another function already named that +PIN_MATRIX_DESCRIPTION = """ +Use the numeric keypad to describe number positions. The layout is: + 7 8 9 + 4 5 6 + 1 2 3 +""".strip() + class TxAPIPSBT(TxApi): def __init__(self, psbt): @@ -100,6 +107,18 @@ def parse_multisig(script): multisig = proto.MultisigRedeemScriptType(m=m, signatures=[b''] * n, pubkeys=pubkeys) return (True, multisig) +# Doesn't init device +class NoInitMixin(ProtocolMixin): + def __init__(self, *args, **kwargs): + super(ProtocolMixin, self).__init__(*args, **kwargs) + self.tx_api = None + +class KeepKey(NoInitMixin, TextUIMixin, BaseClient): + pass + +class KeepKeyDebug(NoInitMixin, DebugLinkMixin, DebugWireMixin, BaseClient): + pass + # This class extends the HardwareWalletClient for Digital Bitbox specific things class KeepkeyClient(HardwareWalletClient): @@ -129,9 +148,15 @@ class KeepkeyClient(HardwareWalletClient): self.password = password os.environ['PASSPHRASE'] = password + def _check_unlocked(self): + self.client.init_device() + if self.client.features.pin_protection and not self.client.features.pin_cached: + raise DeviceNotReadyError('Keepkey is locked. Unlock by using \'promptpin\' and then \'sendpin\'.') + # Must return a dict with the xpub # Retrieves the public key at the specified BIP 32 derivation path def get_pubkey_at_path(self, path): + self._check_unlocked() path = path.replace('h', '\'') path = path.replace('H', '\'') expanded_path = tools.parse_path(path) @@ -144,6 +169,7 @@ class KeepkeyClient(HardwareWalletClient): # Must return a hex string with the signed transaction # The tx must be in the combined unsigned transaction format def sign_tx(self, tx): + self._check_unlocked() # Get this devices master key fingerprint master_key = self.client.get_public_node([0]) @@ -295,6 +321,7 @@ class KeepkeyClient(HardwareWalletClient): # Must return a base64 encoded string with the signed message # The message can be any string def sign_message(self, message, keypath): + self._check_unlocked() keypath = keypath.replace('h', '\'') keypath = keypath.replace('H', '\'') expanded_path = tools.parse_path(keypath) @@ -303,6 +330,7 @@ class KeepkeyClient(HardwareWalletClient): # Display address of specified type on the device. Only supports single-key based addresses. def display_address(self, keypath, p2sh_p2wpkh, bech32): + self._check_unlocked() keypath = keypath.replace('h', '\'') keypath = keypath.replace('H', '\'') expanded_path = tools.parse_path(keypath) @@ -323,6 +351,7 @@ class KeepkeyClient(HardwareWalletClient): # Wipe this device def wipe_device(self): + self._check_unlocked() self.client.wipe_device() return {'success': True} @@ -339,6 +368,33 @@ class KeepkeyClient(HardwareWalletClient): def close(self): self.client.close() + # Prompt for a pin on device + def prompt_pin(self): + self.client.init_device() + if not self.client.features.pin_protection: + raise DeviceAlreadyUnlockedError('This device does not need a PIN') + if self.client.features.pin_cached: + raise DeviceAlreadyUnlockedError('The PIN has already been sent to this device') + print('Use \'sendpin\' to provide the number positions for the PIN as displayed on your device\'s screen', file=sys.stderr) + print(PIN_MATRIX_DESCRIPTION, file=sys.stderr) + self.client.call_raw(messages.Ping(message=b'ping', button_protection=False, pin_protection=True, passphrase_protection=False)) + return {'success': True} + + # Send the pin + def send_pin(self, pin): + if not pin.isdigit(): + raise ValueError("Non-numeric PIN provided") + resp = self.client.call_raw(messages.PinMatrixAck(pin=pin)) + if isinstance(resp, messages.Failure): + self.client.features = self.client.call_raw(messages.GetFeatures()) + if isinstance(self.client.features, messages.Features): + if not self.client.features.pin_protection: + raise DeviceAlreadyUnlockedError('This device does not need a PIN') + if self.client.features.pin_cached: + raise DeviceAlreadyUnlockedError('The PIN has already been sent to this device') + return {'success': False} + return {'success': True} + def enumerate(password=''): results = [] paths = [] @@ -364,6 +420,7 @@ def enumerate(password=''): try: client = KeepkeyClient(path, password) + client.client.init_device() if client.client.features.initialized: master_xpub = client.get_pubkey_at_path('m/0h')['xpub'] d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub) diff --git a/hwilib/devices/ledger.py b/hwilib/devices/ledger.py index ec828f5..17f6cab 100644 --- a/hwilib/devices/ledger.py +++ b/hwilib/devices/ledger.py @@ -268,6 +268,14 @@ class LedgerClient(HardwareWalletClient): def close(self): self.dongle.close() + # Prompt pin + def prompt_pin(self): + raise UnavailableActionError('The Ledger Nano S does not need a PIN sent from the host') + + # Send pin + def send_pin(self): + raise UnavailableActionError('The Ledger Nano S does not need a PIN sent from the host') + def enumerate(password=''): results = [] for d in hid.enumerate(LEDGER_VENDOR_ID, LEDGER_DEVICE_ID): diff --git a/hwilib/devices/trezor.py b/hwilib/devices/trezor.py index 461ce9e..f6e346a 100644 --- a/hwilib/devices/trezor.py +++ b/hwilib/devices/trezor.py @@ -1,10 +1,10 @@ # Trezor interaction script -from ..hwwclient import HardwareWalletClient, DeviceAlreadyInitError, UnavailableActionError +from ..hwwclient import HardwareWalletClient, DeviceAlreadyInitError, DeviceAlreadyUnlockedError, UnavailableActionError, DeviceNotReadyError from trezorlib.client import TrezorClient as Trezor from trezorlib.debuglink import TrezorClientDebugLink from trezorlib.transport import enumerate_devices, get_transport -from trezorlib.ui import ClickUI, mnemonic_words +from trezorlib.ui import ClickUI, mnemonic_words, PIN_MATRIX_DESCRIPTION from trezorlib import protobuf, tools, btc, device from trezorlib import messages as proto from ..base58 import get_xpub_fingerprint, decode, to_address, xpub_main_2_test, get_xpub_fingerprint_hex @@ -15,6 +15,7 @@ import base64 import binascii import json import logging +import sys import os py_enumerate = enumerate # Need to use the enumerate built-in but there's another function already named that @@ -53,6 +54,17 @@ def parse_multisig(script): multisig = proto.MultisigRedeemScriptType(m=m, signatures=[b''] * n, pubkeys=pubkeys) return (True, multisig) +class TrezorNoInit(Trezor): + def __init__(self, transport, ui=None, state=None): + self.transport = transport + self.ui = ui + self.state = state + + if ui is None: + warnings.warn("UI class not supplied. This will probably crash soon.") + + self.session_counter = 0 + # This class extends the HardwareWalletClient for Trezor specific things class TrezorClient(HardwareWalletClient): @@ -63,7 +75,7 @@ class TrezorClient(HardwareWalletClient): transport = get_transport(path) self.client = TrezorClientDebugLink(transport=transport) else: - self.client = Trezor(transport=get_transport(path), ui=ClickUI()) + self.client = TrezorNoInit(transport=get_transport(path), ui=ClickUI()) # if it wasn't able to find a client, throw an error if not self.client: @@ -71,10 +83,17 @@ class TrezorClient(HardwareWalletClient): self.password = password os.environ['PASSPHRASE'] = password + self.client.open() + + def _check_unlocked(self): + self.client.init_device() + if self.client.features.pin_protection and not self.client.features.pin_cached: + raise DeviceNotReadyError('Trezor is locked. Unlock by using \'promptpin\' and then \'sendpin\'.') # Must return a dict with the xpub # Retrieves the public key at the specified BIP 32 derivation path def get_pubkey_at_path(self, path): + self._check_unlocked() expanded_path = tools.parse_path(path) output = btc.get_public_node(self.client, expanded_path) if self.is_testnet: @@ -85,6 +104,7 @@ class TrezorClient(HardwareWalletClient): # Must return a hex string with the signed transaction # The tx must be in the psbt format def sign_tx(self, tx): + self._check_unlocked() # Get this devices master key fingerprint master_key = btc.get_public_node(self.client, [0]) @@ -262,12 +282,14 @@ class TrezorClient(HardwareWalletClient): # Must return a base64 encoded string with the signed message # The message can be any string def sign_message(self, message, keypath): + self._check_unlocked() path = tools.parse_path(keypath) result = btc.sign_message(self.client, 'Bitcoin', path, message) return {'signature': base64.b64encode(result.signature).decode('utf-8')} # Display address of specified type on the device. Only supports single-key based addresses. def display_address(self, keypath, p2sh_p2wpkh, bech32): + self._check_unlocked() expanded_path = tools.parse_path(keypath) address = btc.get_address( self.client, @@ -290,6 +312,7 @@ class TrezorClient(HardwareWalletClient): # Wipe this device def wipe_device(self): + self._check_unlocked() device.wipe(self.client) return {'success': True} @@ -307,6 +330,33 @@ class TrezorClient(HardwareWalletClient): def close(self): self.client.close() + # Prompt for a pin on device + def prompt_pin(self): + self.client.init_device() + if not self.client.features.pin_protection: + raise DeviceAlreadyUnlockedError('This device does not need a PIN') + if self.client.features.pin_cached: + raise DeviceAlreadyUnlockedError('The PIN has already been sent to this device') + print('Use \'sendpin\' to provide the number positions for the PIN as displayed on your device\'s screen', file=sys.stderr) + print(PIN_MATRIX_DESCRIPTION, file=sys.stderr) + self.client.call_raw(proto.Ping(message=b'ping', button_protection=False, pin_protection=True, passphrase_protection=False)) + return {'success': True} + + # Send the pin + def send_pin(self, pin): + self.client.features = self.client.call_raw(proto.GetFeatures()) + if isinstance(self.client.features, proto.Features): + if not self.client.features.pin_protection: + raise DeviceAlreadyUnlockedError('This device does not need a PIN') + if self.client.features.pin_cached: + raise DeviceAlreadyUnlockedError('The PIN has already been sent to this device') + if not pin.isdigit(): + raise ValueError("Non-numeric PIN provided") + resp = self.client.call_raw(proto.PinMatrixAck(pin=pin)) + if isinstance(resp, proto.Failure): + return {'success': False} + return {'success': True} + def enumerate(password=''): results = [] for dev in enumerate_devices(): @@ -317,6 +367,7 @@ def enumerate(password=''): try: client = TrezorClient(d_data['path'], password) + client.client.init_device() if client.client.features.initialized: master_xpub = client.get_pubkey_at_path('m/0h')['xpub'] d_data['fingerprint'] = get_xpub_fingerprint_hex(master_xpub) diff --git a/hwilib/hwwclient.py b/hwilib/hwwclient.py index 02852d2..c7ae2ad 100644 --- a/hwilib/hwwclient.py +++ b/hwilib/hwwclient.py @@ -54,6 +54,14 @@ class HardwareWalletClient(object): raise NotImplementedError('The HardwareWalletClient base class does not ' 'implement this method') + # Prompt pin + def prompt_pin(self): + raise NotImplementedError('The HardwareWalletClient base class does not implement this method') + + # Send pin + def send_pin(self): + raise NotImplementedError('The HardwareWalletClient base class does not implement this method') + class NoPasswordError(Exception): def __init__(self,*args,**kwargs): Exception.__init__(self,*args,**kwargs) @@ -65,3 +73,11 @@ class UnavailableActionError(Exception): class DeviceAlreadyInitError(Exception): def __init__(self,*args,**kwargs): Exception.__init__(self,*args,**kwargs) + +class DeviceNotReadyError(Exception): + def __init__(self,*args,**kwargs): + Exception.__init__(self,*args,**kwargs) + +class DeviceAlreadyUnlockedError(Exception): + def __init__(self,*args,**kwargs): + Exception.__init__(self,*args,**kwargs)