coldcard/HWI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Tests and fixes for remaining Trezor commands

Browse Source 
Tests and fixes for Trezor signmessage, setup, wipe, backup, promptpin, and sendpin
This commit is contained in:
Andrew Chow 2019-02-04 17:47:46 -05:00
parent 40abbc2927
commit cf152b709a
2 changed files with 122 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
hwilib/devices/trezor.py
View File

@ -3,7 +3,7 @@
from ..hwwclient import HardwareWalletClient
from ..errors import ActionCanceledError, BadArgumentError, DeviceAlreadyInitError, DeviceAlreadyUnlockedError, DeviceConnectionError, UnavailableActionError, DeviceNotReadyError
from trezorlib.client import TrezorClient as Trezor
from trezorlib.debuglink import TrezorClientDebugLink
from trezorlib.debuglink import DebugLink, DebugUI, TrezorClientDebugLink
from trezorlib.exceptions import Cancelled
from trezorlib.transport import enumerate_devices, get_transport
from trezorlib.ui import ClickUI, mnemonic_words, PIN_MATRIX_DESCRIPTION
@ -58,15 +58,24 @@ def parse_multisig(script):
return (True, multisig)
class TrezorNoInit(Trezor):
def __init__(self, transport, ui=None, state=None):
self.transport = transport
self.ui = ui
self.state = state
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if ui is None:
warnings.warn("UI class not supplied. This will probably crash soon.")
def init_device(self):
pass
self.session_counter = 0
def actual_init_device(self):
return super().init_device()
class TrezorDebugNoInit(TrezorClientDebugLink):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def init_device(self):
pass
def actual_init_device(self):
return super().init_device()
def trezor_exception(f):
def func(*args, **kwargs):
@ -88,7 +97,7 @@ class TrezorClient(HardwareWalletClient):
if path.startswith('udp'):
logging.debug('Simulator found, using DebugLink')
transport = get_transport(path)
self.client = TrezorClientDebugLink(transport=transport)
self.client = TrezorDebugNoInit(transport=transport)
else:
self.client = TrezorNoInit(transport=get_transport(path), ui=ClickUI())
@ -101,7 +110,7 @@ class TrezorClient(HardwareWalletClient):
self.client.open()
def _check_unlocked(self):
self.client.init_device()
self.client.actual_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\'.')
@ -325,7 +334,7 @@ class TrezorClient(HardwareWalletClient):
# Setup a new device
@trezor_exception
def setup_device(self, label='', passphrase=''):
self.client.init_device()
self.client.actual_init_device()
if self.client.features.initialized:
raise DeviceAlreadyInitError('Device is already initialized. Use wipe first and try again')
passphrase_enabled = False
@ -360,7 +369,7 @@ class TrezorClient(HardwareWalletClient):
# Prompt for a pin on device
@trezor_exception
def prompt_pin(self):
self.client.init_device()
self.client.actual_init_device()
if not self.client.features.pin_protection:
raise DeviceAlreadyUnlockedError('This device does not need a PIN')
if self.client.features.pin_cached:
@ -397,7 +406,7 @@ def enumerate(password=''):
client = None
try:
client = TrezorClient(d_data['path'], password)
client.client.init_device()
client.client.actual_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)

103
test/test_trezor.py
View File

@ -13,11 +13,20 @@ import unittest
from bitcoinrpc.authproxy import AuthServiceProxy, JSONRPCException
from trezorlib.transport import enumerate_devices
from trezorlib.transport.udp import UdpTransport
from trezorlib.debuglink import TrezorClientDebugLink, load_device_by_mnemonic, load_device_by_xprv
from trezorlib import device
from test_device import DeviceEmulator, DeviceTestCase, start_bitcoind, TestDeviceConnect, TestDisplayAddress, TestGetKeypool, TestSignTx
from trezorlib.debuglink import DebugUI, TrezorClientDebugLink, load_device_by_mnemonic, load_device_by_xprv
from trezorlib import device, messages
from test_device import DeviceEmulator, DeviceTestCase, start_bitcoind, TestDeviceConnect, TestDisplayAddress, TestGetKeypool, TestSignMessage, TestSignTx
from hwilib.cli import process_commands
from hwilib.devices.trezor import TrezorClient
from types import MethodType
def get_pin(self, code=None):
if self.pin:
return self.debuglink.encode_pin(self.pin)
else:
return self.debuglink.read_pin_encoded()
class TrezorEmulator(DeviceEmulator):
def __init__(self, path):
@ -102,6 +111,92 @@ class TestTrezorGetxpub(TrezorTestCase):
gxp_res = process_commands(['-t', 'trezor', '-d', 'udp:127.0.0.1:21324', 'getxpub', path_vec['path']])
self.assertEqual(gxp_res['xpub'], path_vec['xpub'])
# Trezor specific management (setup, wipe, restore, backup, promptpin, sendpin) command tests
class TestTrezorManCommands(TrezorTestCase):
def setUp(self):
self.client = self.emulator.start()
self.dev_args = ['-t', 'trezor', '-d', 'udp:127.0.0.1:21324']
def tearDown(self):
self.emulator.stop()
def test_setup_wipe(self):
# Device is init, setup should fail
result = process_commands(self.dev_args + ['setup'])
self.assertEquals(result['code'], -10)
self.assertEquals(result['error'], 'Device is already initialized. Use wipe first and try again')
# Wipe
result = process_commands(self.dev_args + ['wipe'])
self.assertTrue(result['success'])
# Setup
t_client = TrezorClient('udp:127.0.0.1:21324', 'test')
t_client.client.ui.get_pin = MethodType(get_pin, t_client.client.ui)
t_client.client.ui.pin = '1234'
result = t_client.setup_device()
self.assertTrue(result['success'])
# Make sure device is init, setup should fail
result = process_commands(self.dev_args + ['setup'])
self.assertEquals(result['code'], -10)
self.assertEquals(result['error'], 'Device is already initialized. Use wipe first and try again')
def test_backup(self):
result = process_commands(self.dev_args + ['backup'])
self.assertIn('error', result)
self.assertIn('code', result)
self.assertEqual(result['error'], 'The Trezor does not support creating a backup via software')
self.assertEqual(result['code'], -9)
def test_pins(self):
# There's no PIN
result = process_commands(self.dev_args + ['--debug', 'promptpin'])
self.assertEqual(result['error'], 'This device does not need a PIN')
self.assertEqual(result['code'], -11)
result = process_commands(self.dev_args + ['sendpin', '1234'])
self.assertEqual(result['error'], 'This device does not need a PIN')
self.assertEqual(result['code'], -11)
# Set a PIN
device.wipe(self.client)
load_device_by_mnemonic(client=self.client, mnemonic='alcohol woman abuse must during monitor noble actual mixed trade anger aisle', pin='1234', passphrase_protection=False, label='test')
self.client.call(messages.ClearSession())
result = process_commands(self.dev_args + ['promptpin'])
self.assertTrue(result['success'])
# Invalid pins
result = process_commands(self.dev_args + ['sendpin', 'notnum'])
self.assertEqual(result['error'], 'Non-numeric PIN provided')
self.assertEqual(result['code'], -7)
result = process_commands(self.dev_args + ['sendpin', '00000'])
self.assertFalse(result['success'])
# Make sure we get a needs pin message
result = process_commands(self.dev_args + ['getxpub', 'm/0h'])
self.assertEqual(result['code'], -12)
self.assertEqual(result['error'], 'Trezor is locked. Unlock by using \'promptpin\' and then \'sendpin\'.')
# Prompt pin
self.client.call(messages.ClearSession())
result = process_commands(self.dev_args + ['promptpin'])
self.assertTrue(result['success'])
# Send the PIN
self.client.open()
pin = self.client.debug.encode_pin('1234')
result = process_commands(self.dev_args + ['sendpin', pin])
self.assertTrue(result['success'])
# Sending PIN after unlock
result = process_commands(self.dev_args + ['promptpin'])
self.assertEqual(result['error'], 'The PIN has already been sent to this device')
self.assertEqual(result['code'], -11)
result = process_commands(self.dev_args + ['sendpin', '1234'])
self.assertEqual(result['error'], 'The PIN has already been sent to this device')
self.assertEqual(result['code'], -11)
def trezor_test_suite(emulator, rpc, userpass):
# Redirect stderr to /dev/null as it's super spammy
sys.stderr = open(os.devnull, 'w')
@ -119,7 +214,9 @@ def trezor_test_suite(emulator, rpc, userpass):
suite.addTest(DeviceTestCase.parameterize(TestGetKeypool, rpc, userpass, type, path, fingerprint, master_xpub, emulator=dev_emulator))
suite.addTest(DeviceTestCase.parameterize(TestSignTx, rpc, userpass, type, path, fingerprint, master_xpub, emulator=dev_emulator))
suite.addTest(DeviceTestCase.parameterize(TestDisplayAddress, rpc, userpass, type, path, fingerprint, master_xpub, emulator=dev_emulator))
suite.addTest(DeviceTestCase.parameterize(TestSignMessage, rpc, userpass, type, path, fingerprint, master_xpub, emulator=dev_emulator))
suite.addTest(TrezorTestCase.parameterize(TestTrezorGetxpub, emulator=dev_emulator))
suite.addTest(TrezorTestCase.parameterize(TestTrezorManCommands, emulator=dev_emulator))
return suite
if __name__ == '__main__':