diff --git a/cmd.py b/cmd.py index ff632e414..6381dfe03 100755 --- a/cmd.py +++ b/cmd.py @@ -113,6 +113,10 @@ class Commands(object): def wipe_device(self, args): return self.client.wipe_device() + def recovery_device(self, args): + return self.client.recovery_device(args.words, args.passphrase_protection, + args.pin_protection, args.label, 'english') + def load_device(self, args): if not args.mnemonic and not args.xprv: raise Exception("Please provide mnemonic or xprv") @@ -127,7 +131,8 @@ class Commands(object): args.passphrase_protection, args.label, 'english') def reset_device(self, args): - return self.client.reset_device(True, args.strength, args.passphrase, args.pin, args.label, 'english') + return self.client.reset_device(True, args.strength, args.passphrase_protection, + args.pin_protection, args.label, 'english') def sign_message(self, args): return pb2json(self.client.sign_message(args.n, args.message), {'message': args.message}) @@ -155,6 +160,7 @@ class Commands(object): change_pin.help = 'Change new PIN or remove existing' list_coins.help = 'List all supported coin types by the device' wipe_device.help = 'Reset device to factory defaults and remove all private data.' + recovery_device.help = 'Start safe recovery workflow' load_device.help = 'Load custom configuration to the device' reset_device.help = 'Perform device setup and generate new seed' sign_message.help = 'Sign message using address of given path' @@ -190,6 +196,13 @@ class Commands(object): wipe_device.arguments = () + recovery_device.arguments = ( + (('-w', '--words'), {'type': int}), + (('-p', '--pin-protection'), {'action': 'store_true', 'default': False}), + (('-r', '--passphrase-protection'), {'action': 'store_true', 'default': False}), + (('-l', '--label'), {'type': str, 'default': ''}), + ) + load_device.arguments = ( (('-m', '--mnemonic'), {'type': str, 'nargs': '+'}), (('-x', '--xprv'), {'type': str}), @@ -200,8 +213,8 @@ class Commands(object): reset_device.arguments = ( (('-t', '--strength'), {'type': int, 'choices': [128, 192, 256], 'default': 128}), - (('-p', '--pin'), {'action': 'store_true', 'default': False}), - (('-r', '--passphrase'), {'action': 'store_true', 'default': False}), + (('-p', '--pin-protection'), {'action': 'store_true', 'default': False}), + (('-r', '--passphrase-protection'), {'action': 'store_true', 'default': False}), (('-l', '--label'), {'type': str, 'default': ''}), ) diff --git a/trezorlib/client.py b/trezorlib/client.py index 9e8d90215..97dfcebe7 100755 --- a/trezorlib/client.py +++ b/trezorlib/client.py @@ -26,6 +26,9 @@ def pin_func(input_text, message=None): def passphrase_func(input_text): return show_input(input_text) +def word_func(): + return raw_input("Enter one word of mnemonic: ") + class CallException(Exception): pass @@ -39,7 +42,7 @@ class TrezorClient(object): def __init__(self, transport, debuglink=None, message_func=show_message, input_func=show_input, pin_func=pin_func, passphrase_func=passphrase_func, - blockchain_api=None, debug=False): + word_func=word_func, blockchain_api=None, debug=False): self.transport = transport self.debuglink = debuglink @@ -47,6 +50,7 @@ class TrezorClient(object): self.input_func = input_func self.pin_func = pin_func self.passphrase_func = passphrase_func + self.word_func = word_func self.debug = debug @@ -325,18 +329,38 @@ class TrezorClient(object): self.init_device() return ret + def recovery_device(self, word_count, passphrase_protection, pin_protection, label, language): + if word_count not in (12, 18, 24): + raise Exception("Invalid word count. Use 12/18/24") + + res = self.call(proto.RecoveryDevice(word_count=int(word_count), + passphrase_protection=bool(passphrase_protection), + pin_protection=bool(pin_protection), + label=label, + language=language)) + + while isinstance(res, proto.WordRequest): + word = self.word_func() + res = self.call(proto.WordAck(word=word)) + + if not isinstance(res, proto.Success): + raise Exception("Recovery device failed") + + self.init_device() + return True + def reset_device(self, display_random, strength, passphrase_protection, pin_protection, label, language): if self.features.initialized: raise Exception("Device is initialized already. Call wipe_device() and try again.") # Begin with device reset workflow msg = proto.ResetDevice(display_random=display_random, - strength=strength, - language=language, - passphrase_protection=bool(passphrase_protection), - pin_protection=bool(pin_protection), - label=label - ) + strength=strength, + language=language, + passphrase_protection=bool(passphrase_protection), + pin_protection=bool(pin_protection), + label=label) + resp = self.call(msg) if not isinstance(resp, proto.EntropyRequest): raise Exception("Invalid response, expected EntropyRequest")