From 3a735d40acecbcee689e13e5d6dbef9d10c01584 Mon Sep 17 00:00:00 2001
From: Neil Booth <kyuupichan@gmail.com>
Date: Sun, 24 Jan 2016 10:39:04 +0900
Subject: [PATCH] Keepkey/Trezor: fix 2 issues

1: In KeepKey recovery we let the user type a fifth
   letter which causes the firmware to abort the
   protocol.  Only allow 4 letters.
2: Handle exceptions thrown during device initialization
   properly
---
 plugins/trezor/plugin.py     | 13 ++++++++-----
 plugins/trezor/qt_generic.py | 29 +++++++++++++++++++++--------
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/plugins/trezor/plugin.py b/plugins/trezor/plugin.py
index 520b69fab..60b9b5b5c 100644
--- a/plugins/trezor/plugin.py
+++ b/plugins/trezor/plugin.py
@@ -323,18 +323,21 @@ class TrezorCompatiblePlugin(BasePlugin, ThreadJob):
                 pin = pin_protection  # It's the pin, not a boolean
                 client.load_device_by_xprv(item, pin, passphrase_protection,
                                            label, language)
+            # After successful initialization create accounts
+            wallet.create_hd_account(None)
 
         return initialize_method
 
-    def setup_device(self, wallet, on_done):
+    def setup_device(self, wallet, on_done, on_error):
         '''Called when creating a new wallet.  Select the device to use.  If
         the device is uninitialized, go through the intialization
         process.  Then create the wallet accounts.'''
-        create_hd_task = partial(wallet.create_hd_account, None)
         initialized = self.select_device(wallet)
-        if not initialized:
-            wallet.thread.add(self.initialize_device(wallet))
-        wallet.thread.add(create_hd_task, on_done=on_done)
+        if initialized:
+            task = partial(wallet.create_hd_account, None)
+        else:
+            task = self.initialize_device(wallet)
+        wallet.thread.add(task, on_done=on_done, on_error=on_error)
 
     def unpaired_devices(self, handler):
         '''Returns all connected, unpaired devices as a list of clients and a
diff --git a/plugins/trezor/qt_generic.py b/plugins/trezor/qt_generic.py
index c3a4b45f8..caef3a2bf 100644
--- a/plugins/trezor/qt_generic.py
+++ b/plugins/trezor/qt_generic.py
@@ -29,8 +29,8 @@ PASSPHRASE_NOT_PIN = _(
     "Only change this if you are sure you understand it.")
 CHARACTER_RECOVERY = (
     "Use the recovery cipher shown on your device to input your seed words.  "
-    "The cipher updates with every letter.  After at most 4 letters the "
-    "device will auto-complete each word.\n"
+    "The cipher changes with every keypress.\n"
+    "After at most 4 letters the device will auto-complete a word.\n"
     "Press SPACE or the Accept Word button to accept the device's auto-"
     "completed word and advance to the next one.\n"
     "Press BACKSPACE to go back a character or word.\n"
@@ -93,17 +93,24 @@ class CharacterDialog(WindowModalDialog):
             if n == self.character_pos:
                 button.setFocus()
 
+    def is_valid_alpha_space(self, key):
+        # Auto-completion requires at least 3 characters
+        if key == ord(' ') and self.character_pos >= 3:
+            return True
+        # Firmware aborts protocol if the 5th character is non-space
+        if self.character_pos >= 4:
+            return False
+        return (key >= ord('a') and key <= ord('z')
+                or (key >= ord('A') and key <= ord('Z')))
+
     def process_key(self, key):
         self.data = None
         if key == Qt.Key_Return and self.finished_button.isEnabled():
             self.data = {'done': True}
         elif key == Qt.Key_Backspace and (self.word_pos or self.character_pos):
             self.data = {'delete': True}
-        elif ((key >= ord('a') and key <= ord('z'))
-              or (key >= ord('A') and key <= ord('Z'))
-              or (key == ord(' ') and self.character_pos >= 3)):
-            char = chr(key).lower()
-            self.data = {'character': char}
+        elif self.is_valid_alpha_space(key):
+            self.data = {'character': chr(key).lower()}
         if self.data:
             self.loop.exit(0)
 
@@ -368,8 +375,14 @@ def qt_plugin_class(base_plugin_class):
         wallet.thread = TaskThread(wizard, wizard.on_error)
         # Setup device and create accounts in separate thread; wait until done
         loop = QEventLoop()
-        self.setup_device(wallet, loop.quit)
+        exc_info = []
+        self.setup_device(wallet, on_done=loop.quit,
+                          on_error=lambda info: exc_info.extend(info))
         loop.exec_()
+        # If an exception was thrown, show to user and exit install wizard
+        if exc_info:
+            wizard.on_error(exc_info)
+            raise UserCancelled
 
     @hook
     def receive_menu(self, menu, addrs, wallet):