Merge pull request #4329 from SomberNight/trezor_matrix

Trezor: Matrix recovery support

plugins/trezor/clientbase.py

@@ -86,6 +86,15 @@ class GuiMixin(object):
         return self.proto.PassphraseStateAck()
 
     def callback_WordRequest(self, msg):
+        if (msg.type is not None
+            and msg.type in (self.types.WordRequestType.Matrix9,
+                             self.types.WordRequestType.Matrix6)):
+            num = 9 if msg.type == self.types.WordRequestType.Matrix9 else 6
+            char = self.handler.get_matrix(num)
+            if char == 'x':
+                return self.proto.Cancel()
+            return self.proto.WordAck(word=char)
+
         self.step += 1
         msg = _("Step {}/24.  Enter seed word as explained on "
                 "your {}:").format(self.step, self.device)
@@ -226,6 +235,10 @@ class TrezorClientBase(GuiMixin, PrintError):
     def atleast_version(self, major, minor=0, patch=0):
         return self.firmware_version() >= (major, minor, patch)
 
+    def get_trezor_model(self):
+        """Returns '1' for Trezor One, 'T' for Trezor T."""
+        return self.features.model
+
     @staticmethod
     def wrapper(func):
         '''Wrap methods to clear any message box they opened.'''

plugins/trezor/qt.py

@@ -12,7 +12,8 @@ from electrum.util import PrintError, UserCancelled, bh2u
 from electrum.wallet import Wallet, Standard_Wallet
 
 from ..hw_wallet.qt import QtHandlerBase, QtPluginBase
-from .trezor import TrezorPlugin, TIM_NEW, TIM_RECOVER, TIM_MNEMONIC
+from .trezor import (TrezorPlugin, TIM_NEW, TIM_RECOVER, TIM_MNEMONIC,
+                     RECOVERY_TYPE_SCRAMBLED_WORDS, RECOVERY_TYPE_MATRIX)
 
 
 PASSPHRASE_HELP_SHORT =_(
@@ -30,16 +31,87 @@ PASSPHRASE_NOT_PIN = _(
     "If you forget a passphrase you will be unable to access any "
     "bitcoins in the wallet behind it.  A passphrase is not a PIN. "
     "Only change this if you are sure you understand it.")
+MATRIX_RECOVERY = _(
+    "Enter the recovery words by pressing the buttons according to what "
+    "the device shows on its display.  You can also use your NUMPAD.\n"
+    "Press BACKSPACE to go back a choice or word.\n")
+
+
+class MatrixDialog(WindowModalDialog):
+
+    def __init__(self, parent):
+        super(MatrixDialog, self).__init__(parent)
+        self.setWindowTitle(_("Trezor Matrix Recovery"))
+        self.num = 9
+        self.loop = QEventLoop()
+
+        vbox = QVBoxLayout(self)
+        vbox.addWidget(WWLabel(MATRIX_RECOVERY))
+
+        grid = QGridLayout()
+        grid.setSpacing(0)
+        self.char_buttons = []
+        for y in range(3):
+            for x in range(3):
+                button = QPushButton('?')
+                button.clicked.connect(partial(self.process_key, ord('1') + y * 3 + x))
+                grid.addWidget(button, 3 - y, x)
+                self.char_buttons.append(button)
+        vbox.addLayout(grid)
+
+        self.backspace_button = QPushButton("<=")
+        self.backspace_button.clicked.connect(partial(self.process_key, Qt.Key_Backspace))
+        self.cancel_button = QPushButton(_("Cancel"))
+        self.cancel_button.clicked.connect(partial(self.process_key, Qt.Key_Escape))
+        buttons = Buttons(self.backspace_button, self.cancel_button)
+        vbox.addSpacing(40)
+        vbox.addLayout(buttons)
+        self.refresh()
+        self.show()
+
+    def refresh(self):
+        for y in range(3):
+            self.char_buttons[3 * y + 1].setEnabled(self.num == 9)
+
+    def is_valid(self, key):
+        return key >= ord('1') and key <= ord('9')
+
+    def process_key(self, key):
+        self.data = None
+        if key == Qt.Key_Backspace:
+            self.data = '\010'
+        elif key == Qt.Key_Escape:
+            self.data = 'x'
+        elif self.is_valid(key):
+            self.char_buttons[key - ord('1')].setFocus()
+            self.data = '%c' % key
+        if self.data:
+            self.loop.exit(0)
+
+    def keyPressEvent(self, event):
+        self.process_key(event.key())
+        if not self.data:
+            QDialog.keyPressEvent(self, event)
+
+    def get_matrix(self, num):
+        self.num = num
+        self.refresh()
+        self.loop.exec_()
 
 
 class QtHandler(QtHandlerBase):
 
     pin_signal = pyqtSignal(object)
+    matrix_signal = pyqtSignal(object)
+    close_matrix_dialog_signal = pyqtSignal()
 
     def __init__(self, win, pin_matrix_widget_class, device):
         super(QtHandler, self).__init__(win, device)
         self.pin_signal.connect(self.pin_dialog)
+        self.matrix_signal.connect(self.matrix_recovery_dialog)
+        self.close_matrix_dialog_signal.connect(self._close_matrix_dialog)
         self.pin_matrix_widget_class = pin_matrix_widget_class
+        self.matrix_dialog = None
 
     def get_pin(self, msg):
         self.done.clear()
@@ -47,6 +119,23 @@ class QtHandler(QtHandlerBase):
         self.done.wait()
         return self.response
 
+    def get_matrix(self, msg):
+        self.done.clear()
+        self.matrix_signal.emit(msg)
+        self.done.wait()
+        data = self.matrix_dialog.data
+        if data == 'x':
+            self.close_matrix_dialog()
+        return data
+
+    def _close_matrix_dialog(self):
+        if self.matrix_dialog:
+            self.matrix_dialog.accept()
+            self.matrix_dialog = None
+
+    def close_matrix_dialog(self):
+        self.close_matrix_dialog_signal.emit()
+
     def pin_dialog(self, msg):
         # Needed e.g. when resetting a device
         self.clear_dialog()
@@ -61,6 +150,12 @@ class QtHandler(QtHandlerBase):
         self.response = str(matrix.get_value())
         self.done.set()
 
+    def matrix_recovery_dialog(self, msg):
+        if not self.matrix_dialog:
+            self.matrix_dialog = MatrixDialog(self.top_level_window())
+        self.matrix_dialog.get_matrix(msg)
+        self.done.set()
+
 
 class QtPlugin(QtPluginBase):
     # Derived classes must provide the following class-static variables:
@@ -86,7 +181,7 @@ class QtPlugin(QtPluginBase):
         if device_id:
             SettingsDialog(window, self, keystore, device_id).exec_()
 
-    def request_trezor_init_settings(self, wizard, method, device):
+    def request_trezor_init_settings(self, wizard, method, model):
         vbox = QVBoxLayout()
         next_enabled = True
         label = QLabel(_("Enter a label to name your device:"))
@@ -107,12 +202,12 @@ class QtPlugin(QtPluginBase):
             gb.setLayout(hbox1)
             vbox.addWidget(gb)
             gb.setTitle(_("Select your seed length:"))
-            bg = QButtonGroup()
+            bg_numwords = QButtonGroup()
             for i, count in enumerate([12, 18, 24]):
                 rb = QRadioButton(gb)
                 rb.setText(_("%d words") % count)
-                bg.addButton(rb)
-                bg.setId(rb, i)
+                bg_numwords.addButton(rb)
+                bg_numwords.setId(rb, i)
                 hbox1.addWidget(rb)
                 rb.setChecked(True)
             cb_pin = QCheckBox(_('Enable PIN protection'))
@@ -155,16 +250,42 @@ class QtPlugin(QtPluginBase):
         vbox.addWidget(passphrase_warning)
         vbox.addWidget(cb_phrase)
 
+        # ask for recovery type (random word order OR matrix)
+        if method == TIM_RECOVER and not model == 'T':
+            gb_rectype = QGroupBox()
+            hbox_rectype = QHBoxLayout()
+            gb_rectype.setLayout(hbox_rectype)
+            vbox.addWidget(gb_rectype)
+            gb_rectype.setTitle(_("Select recovery type:"))
+            bg_rectype = QButtonGroup()
+
+            rb1 = QRadioButton(gb_rectype)
+            rb1.setText(_('Scrambled words'))
+            bg_rectype.addButton(rb1)
+            bg_rectype.setId(rb1, RECOVERY_TYPE_SCRAMBLED_WORDS)
+            hbox_rectype.addWidget(rb1)
+            rb1.setChecked(True)
+
+            rb2 = QRadioButton(gb_rectype)
+            rb2.setText(_('Matrix'))
+            bg_rectype.addButton(rb2)
+            bg_rectype.setId(rb2, RECOVERY_TYPE_MATRIX)
+            hbox_rectype.addWidget(rb2)
+        else:
+            bg_rectype = None
+
         wizard.exec_layout(vbox, next_enabled=next_enabled)
 
         if method in [TIM_NEW, TIM_RECOVER]:
-            item = bg.checkedId()
+            item = bg_numwords.checkedId()
             pin = cb_pin.isChecked()
+            recovery_type = bg_rectype.checkedId() if bg_rectype else None
         else:
             item = ' '.join(str(clean_text(text)).split())
             pin = str(pin.text())
+            recovery_type = None
 
-        return (item, name.text(), pin, cb_phrase.isChecked())
+        return (item, name.text(), pin, cb_phrase.isChecked(), recovery_type)
 
 
 class Plugin(TrezorPlugin, QtPlugin):

plugins/trezor/trezor.py

@@ -17,6 +17,7 @@ from ..hw_wallet import HW_PluginBase
 
 # TREZOR initialization methods
 TIM_NEW, TIM_RECOVER, TIM_MNEMONIC, TIM_PRIVKEY = range(0, 4)
+RECOVERY_TYPE_SCRAMBLED_WORDS, RECOVERY_TYPE_MATRIX = range(0, 2)
 
 # script "generation"
 SCRIPT_GEN_LEGACY, SCRIPT_GEN_P2SH_SEGWIT, SCRIPT_GEN_NATIVE_SEGWIT = range(0, 3)
@@ -192,9 +193,12 @@ class TrezorPlugin(HW_PluginBase):
             (TIM_MNEMONIC, _("Upload a BIP39 mnemonic to generate the seed")),
             (TIM_PRIVKEY, _("Upload a master private key"))
         ]
+        devmgr = self.device_manager()
+        client = devmgr.client_by_id(device_id)
+        model = client.get_trezor_model()
         def f(method):
             import threading
-            settings = self.request_trezor_init_settings(wizard, method, self.device)
+            settings = self.request_trezor_init_settings(wizard, method, model)
             t = threading.Thread(target=self._initialize_device_safe, args=(settings, method, device_id, wizard, handler))
             t.setDaemon(True)
             t.start()
@@ -213,9 +217,9 @@ class TrezorPlugin(HW_PluginBase):
             wizard.loop.exit(0)
 
     def _initialize_device(self, settings, method, device_id, wizard, handler):
-        item, label, pin_protection, passphrase_protection = settings
+        item, label, pin_protection, passphrase_protection, recovery_type = settings
 
-        if method == TIM_RECOVER:
+        if method == TIM_RECOVER and recovery_type == RECOVERY_TYPE_SCRAMBLED_WORDS:
             handler.show_error(_(
                 "You will be asked to enter 24 words regardless of your "
                 "seed's actual length.  If you enter a word incorrectly or "
@@ -238,8 +242,15 @@ class TrezorPlugin(HW_PluginBase):
         elif method == TIM_RECOVER:
             word_count = 6 * (item + 2)  # 12, 18 or 24
             client.step = 0
+            if recovery_type == RECOVERY_TYPE_SCRAMBLED_WORDS:
+                recovery_type_trezor = self.types.RecoveryDeviceType.ScrambledWords
+            else:
+                recovery_type_trezor = self.types.RecoveryDeviceType.Matrix
             client.recovery_device(word_count, passphrase_protection,
-                                       pin_protection, label, language)
+                                   pin_protection, label, language,
+                                   type=recovery_type_trezor)
+            if recovery_type == RECOVERY_TYPE_MATRIX:
+                handler.close_matrix_dialog()
         elif method == TIM_MNEMONIC:
             pin = pin_protection  # It's the pin, not a boolean
             client.load_device_by_mnemonic(str(item), pin,