Introduce and implement prompt_pin and send_pin methods
Only the Trezor and KeepKey need to have a PIN entered. However this PIN entering is canceled if the device is re-initialized which happens every time the client is created. To work around this, a client subclass for each is introduced which does not initialize the device on creation. Instead device initialization is done at the beginning of each call which requires it. At that time, a check is also done to ensure that the device has the PIN cached. This allows for send_pin to actually be able to send the PIN after promptpin.
This commit is contained in:
Andrew Chow
parent
bd57b8fdd8
commit
bb88cc5e6f
@ -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):
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user