Skip to content

Commit

Permalink
trezor: Set coin name as testnet when --testnet
Browse files Browse the repository at this point in the history
Also update displayaddress test to use --testnet
  • Loading branch information
achow101 committed Jun 6, 2020
1 parent d662fad commit f862135
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 10 deletions.
20 changes: 10 additions & 10 deletions hwilib/devices/trezor.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ def __init__(self, path, password='', expert=False):
self.type = 'Trezor'

def _check_unlocked(self):
self.coin_name = 'Testnet' if self.is_testnet else 'Bitcoin'
self.client.init_device()
if self.client.features.model == 'T':
self.client.ui.disallow_passphrase()
Expand All @@ -124,7 +125,7 @@ def get_pubkey_at_path(self, path):
expanded_path = tools.parse_path(path)
except ValueError as e:
raise BadArgumentError(str(e))
output = btc.get_public_node(self.client, expanded_path)
output = btc.get_public_node(self.client, expanded_path, coin_name=self.coin_name)
if self.is_testnet:
result = {'xpub': xpub_main_2_test(output.xpub)}
else:
Expand All @@ -142,7 +143,7 @@ def sign_tx(self, tx):
self._check_unlocked()

# Get this devices master key fingerprint
master_key = btc.get_public_node(self.client, [0])
master_key = btc.get_public_node(self.client, [0x80000000], coin_name='Bitcoin')
master_fp = get_xpub_fingerprint(master_key.xpub)

# Do multiple passes for multisig
Expand Down Expand Up @@ -321,10 +322,7 @@ def ignore_input():
tx_details = proto.SignTx()
tx_details.version = tx.tx.nVersion
tx_details.lock_time = tx.tx.nLockTime
if self.is_testnet:
signed_tx = btc.sign_tx(self.client, "Testnet", inputs, outputs, tx_details, prevtxs)
else:
signed_tx = btc.sign_tx(self.client, "Bitcoin", inputs, outputs, tx_details, prevtxs)
signed_tx = btc.sign_tx(self.client, self.coin_name, inputs, outputs, tx_details, prevtxs)

# Each input has one signature
for input_num, (psbt_in, sig) in py_enumerate(list(zip(tx.inputs, signed_tx[0]))):
Expand All @@ -346,20 +344,21 @@ def ignore_input():
def sign_message(self, message, keypath):
self._check_unlocked()
path = tools.parse_path(keypath)
result = btc.sign_message(self.client, 'Bitcoin', path, message)
result = btc.sign_message(self.client, self.coin_name, 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.
@trezor_exception
def display_address(self, keypath, p2sh_p2wpkh, bech32):
self._check_unlocked()
expanded_path = tools.parse_path(keypath)
script_type = proto.InputScriptType.SPENDWITNESS if bech32 else (proto.InputScriptType.SPENDP2SHWITNESS if p2sh_p2wpkh else proto.InputScriptType.SPENDADDRESS)
address = btc.get_address(
self.client,
"Testnet" if self.is_testnet else "Bitcoin",
self.coin_name,
expanded_path,
show_display=True,
script_type=proto.InputScriptType.SPENDWITNESS if bech32 else (proto.InputScriptType.SPENDP2SHWITNESS if p2sh_p2wpkh else proto.InputScriptType.SPENDADDRESS)
script_type=script_type
)
return {'address': address}

Expand Down Expand Up @@ -406,6 +405,7 @@ def close(self):
# Prompt for a pin on device
@trezor_exception
def prompt_pin(self):
self.coin_name = 'Testnet' if self.is_testnet else 'Bitcoin'
self.client.open()
self.client.init_device()
if not self.client.features.pin_protection:
Expand All @@ -414,7 +414,7 @@ def prompt_pin(self):
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.GetPublicKey(address_n=[0x8000002c, 0x80000001, 0x80000000], ecdsa_curve_name=None, show_display=False, coin_name=None, script_type=proto.InputScriptType.SPENDADDRESS))
self.client.call_raw(proto.GetPublicKey(address_n=[0x8000002c, 0x80000001, 0x80000000], ecdsa_curve_name=None, show_display=False, coin_name=self.coin_name, script_type=proto.InputScriptType.SPENDADDRESS))
return {'success': True}

# Send the pin
Expand Down
5 changes: 5 additions & 0 deletions test/test_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -456,6 +456,11 @@ def test_big_tx(self):
pass

class TestDisplayAddress(DeviceTestCase):
def setUp(self):
if '--testnet' not in self.dev_args:
self.dev_args.append('--testnet')
self.emulator.start()

def test_display_address_bad_args(self):
result = self.do_command(self.dev_args + ['displayaddress', '--sh_wpkh', '--wpkh', '--path', 'm/49h/1h/0h/0/0'])
self.assertIn('error', result)
Expand Down

0 comments on commit f862135

Please sign in to comment.