From a2792cda192bcdd57c995fb02528fc0a2c1217b2 Mon Sep 17 00:00:00 2001
From: mruddy <mruddy@users.noreply.github.com>
Date: Sat, 11 Feb 2017 14:15:09 -0500
Subject: [PATCH] load_device_by_mnemonic and recovery_device: expand unique
 mnemonic word prefix matches (#96)

---
 trezorctl           |  7 +++++--
 trezorlib/client.py | 16 ++++++++++++++--
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/trezorctl b/trezorctl
index fa2439f2a2..a28b02a2ba 100755
--- a/trezorctl
+++ b/trezorctl
@@ -271,7 +271,7 @@ class Commands(object):
                     'matrix':    types.RecoveryDeviceType_Matrix }
         return self.client.recovery_device(args.words, args.passphrase_protection,
                                            args.pin_protection, args.label, 'english',
-                                           typemap[args.type])
+                                           typemap[args.type], args.expand)
 
     def load_device(self, args):
         if not args.mnemonic and not args.xprv:
@@ -282,7 +282,8 @@ class Commands(object):
             return self.client.load_device_by_mnemonic(mnemonic, args.pin,
                                                        args.passphrase_protection,
                                                        args.label, 'english',
-                                                       args.skip_checksum)
+                                                       args.skip_checksum,
+                                                       args.expand)
         else:
             return self.client.load_device_by_xprv(args.xprv, args.pin,
                                                    args.passphrase_protection,
@@ -456,6 +457,7 @@ class Commands(object):
 
     recovery_device.arguments = (
         (('-w', '--words'), {'type': int, 'choices': [12, 18, 24], 'default': 24}),
+        (('-e', '--expand'), {'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': ''}),
@@ -464,6 +466,7 @@ class Commands(object):
 
     load_device.arguments = (
         (('-m', '--mnemonic'), {'type': str, 'nargs': '+'}),
+        (('-e', '--expand'), {'action': 'store_true', 'default': False}),
         (('-x', '--xprv'), {'type': str}),
         (('-p', '--pin'), {'type': str, 'default': ''}),
         (('-r', '--passphrase-protection'), {'action': 'store_true', 'default': False}),
diff --git a/trezorlib/client.py b/trezorlib/client.py
index b89059d5b6..7772fab178 100644
--- a/trezorlib/client.py
+++ b/trezorlib/client.py
@@ -277,6 +277,8 @@ class TextUIMixin(object):
             word = raw_input()
         except NameError:
             word = input() # Python 3
+        if self.expand:
+            word = self.mnemonic_wordlist.expand_word(word)
         return proto.WordAck(word=word)
 
 class DebugLinkMixin(object):
@@ -812,7 +814,7 @@ class ProtocolMixin(object):
 
     @field('message')
     @expect(proto.Success)
-    def recovery_device(self, word_count, passphrase_protection, pin_protection, label, language, type=types.RecoveryDeviceType_ScrambledWords):
+    def recovery_device(self, word_count, passphrase_protection, pin_protection, label, language, type=types.RecoveryDeviceType_ScrambledWords, expand=False):
         if self.features.initialized:
             raise Exception("Device is initialized already. Call wipe_device() and try again.")
 
@@ -820,6 +822,12 @@ class ProtocolMixin(object):
             raise Exception("Invalid word count. Use 12/18/24")
 
         self.recovery_matrix_first_pass = True
+
+        self.expand = expand
+        if self.expand:
+            # optimization to load the wordlist once, instead of for each recovery word
+            self.mnemonic_wordlist = Mnemonic('english')
+
         res = self.call(proto.RecoveryDevice(word_count=int(word_count),
                                    passphrase_protection=bool(passphrase_protection),
                                    pin_protection=bool(pin_protection),
@@ -858,7 +866,7 @@ class ProtocolMixin(object):
 
     @field('message')
     @expect(proto.Success)
-    def load_device_by_mnemonic(self, mnemonic, pin, passphrase_protection, label, language, skip_checksum=False):
+    def load_device_by_mnemonic(self, mnemonic, pin, passphrase_protection, label, language, skip_checksum=False, expand=False):
         # Convert mnemonic to UTF8 NKFD
         mnemonic = Mnemonic.normalize_string(mnemonic)
 
@@ -866,6 +874,10 @@ class ProtocolMixin(object):
         mnemonic = normalize_nfc(mnemonic)
 
         m = Mnemonic('english')
+
+        if expand:
+            mnemonic = m.expand(mnemonic)
+
         if not skip_checksum and not m.check(mnemonic):
             raise Exception("Invalid mnemonic checksum")