Merge pull request #86 from cryptoadvance/fix-mnemonic

allow fixing invalid mnemonic

src/apps/label.py

@@ -26,9 +26,7 @@ async def process_host_command(self, stream, show_screen):
         prefix = self.get_prefix(stream)
         if prefix == b"getlabel":
             label = self.get_label()
-            obj = {
-                "title": "Device's label is: %s" % label,
-            }
+            obj = {"title": "Device's label is: %s" % label}
             stream = BytesIO(label.encode())
             return stream, obj
         elif prefix == b"setlabel":
@@ -43,9 +41,7 @@ async def process_host_command(self, stream, show_screen):
             if res is False:
                 return None
             self.set_label(label)
-            obj = {
-                "title": "New device label: %s" % label,
-            }
+            obj = {"title": "New device label: %s" % label}
             return BytesIO(label.encode()), obj
         else:
             raise AppError("Invalid command")

src/apps/wallets/manager.py

@@ -33,9 +33,7 @@ class WalletManager(BaseApp):
     """
 
     button = "Wallets"
-    WALLETS = [
-        Wallet,
-    ]
+    WALLETS = [Wallet]
 
     def __init__(self, path):
         self.root_path = path

src/apps/wallets/wallet.py

@@ -19,10 +19,7 @@ class Wallet:
     wrapped=True - nested segwit
     """
 
-    SCRIPTS = [
-        SingleKey,
-        Multisig,
-    ]
+    SCRIPTS = [SingleKey, Multisig]
     GAP_LIMIT = 20
 
     def __init__(self, script, wrapped=False, path=None, name="Untitled"):

src/apps/xpubs/xpubs.py

@@ -16,10 +16,7 @@ class XpubApp(BaseApp):
     """
 
     button = "Master public keys"
-    prefixes = [
-        b"fingerprint",
-        b"xpub",
-    ]
+    prefixes = [b"fingerprint", b"xpub"]
 
     def __init__(self, path):
         # we don't need to store anything

src/gui/components/mnemonic.py

@@ -18,10 +18,10 @@ def __init__(self, *args, **kwargs):
 
         self.set_col_cnt(4)
         self.set_row_cnt(12)
-        self.set_col_width(0, 50)
-        self.set_col_width(2, 50)
-        self.set_col_width(1, 150)
-        self.set_col_width(3, 150)
+        self.set_col_width(0, 40)
+        self.set_col_width(2, 40)
+        self.set_col_width(1, 180)
+        self.set_col_width(3, 180)
 
         self.set_style(lv.page.STYLE.BG, cell_style)
         self.set_style(lv.table.STYLE.CELL1, cell_style)
@@ -65,14 +65,19 @@ def del_char(self):
         self.update()
 
     def autocomplete_word(self, word):
+        if len(self.words) > 24:
+            return
         if len(self.words) == 0:
             self.words.append(word)
         else:
             self.words[-1] = word
-        self.words.append("")
+        if len(self.words) < 24:
+            self.words.append("")
         self.update()
 
     def add_char(self, c):
+        if len(self.words) > 24:
+            return
         if len(self.words) == 0:
             self.words.append(c)
         else:

src/gui/screens/mnemonic.py

@@ -61,7 +61,13 @@ def confirm(self):
 
 
 class RecoverMnemonicScreen(MnemonicScreen):
-    def __init__(self, checker=None, lookup=None, title="Enter your recovery phrase"):
+    # button indexes
+    BTN_NEXT = 28
+    BTN_DONE = 29
+
+    def __init__(
+        self, checker=None, lookup=None, fixer=None, title="Enter your recovery phrase"
+    ):
         super().__init__("", title)
         self.table.align(self.title, lv.ALIGN.OUT_BOTTOM_MID, 0, 10)
         self.checker = checker
@@ -115,15 +121,22 @@ def __init__(self, checker=None, lookup=None, title="Enter your recovery phrase"
 
         if lookup is not None:
             # Next word button inactive
-            self.kb.set_btn_ctrl(28, lv.btnm.CTRL.INACTIVE)
+            self.kb.set_btn_ctrl(self.BTN_NEXT, lv.btnm.CTRL.INACTIVE)
         if checker is not None:
             # Done inactive
-            self.kb.set_btn_ctrl(29, lv.btnm.CTRL.INACTIVE)
+            self.kb.set_btn_ctrl(self.BTN_DONE, lv.btnm.CTRL.INACTIVE)
         self.kb.set_width(HOR_RES)
-        self.kb.set_height(VER_RES // 3)
+        self.kb.set_height(260)
         self.kb.align(self, lv.ALIGN.IN_BOTTOM_MID, 0, 0)
         self.kb.set_event_cb(self.callback)
 
+        self.fixer = fixer
+        if fixer is not None:
+            self.fix_button = add_button("fix", on_release(self.fix_cb), self)
+            self.fix_button.set_size(55, 30)
+            # position it out of the screen but on correct y
+            self.fix_button.align(self.table, lv.ALIGN.OUT_BOTTOM_MID, -400, -38)
+
         if lookup is not None:
             self.autocomplete.set_width(HOR_RES)
             self.autocomplete.set_height(50)
@@ -132,6 +145,10 @@ def __init__(self, checker=None, lookup=None, title="Enter your recovery phrase"
             self.autocomplete.set_map(words)
             self.autocomplete.set_event_cb(self.select_word)
 
+    def fix_cb(self):
+        self.table.set_mnemonic(self.fixer(self.table.get_mnemonic()))
+        self.check_buttons()
+
     def select_word(self, obj, event):
         if event != lv.EVENT.RELEASED:
             return
@@ -150,22 +167,43 @@ def check_buttons(self):
         mnemonic = self.table.get_mnemonic()
         # check if we can autocomplete the last word
         if self.lookup is not None:
-            self.kb.set_btn_ctrl(28, lv.btnm.CTRL.INACTIVE)
+            self.kb.set_btn_ctrl(self.BTN_NEXT, lv.btnm.CTRL.INACTIVE)
             word = self.table.get_last_word()
             candidates = self.lookup(word, 4)
             self.autocomplete.set_map(candidates + [""])
             if len(candidates) == 1 or word in candidates:
-                self.kb.clear_btn_ctrl(28, lv.btnm.CTRL.INACTIVE)
+                self.kb.clear_btn_ctrl(self.BTN_NEXT, lv.btnm.CTRL.INACTIVE)
                 if len(candidates) == 1:
                     mnemonic = " ".join(self.table.words[:-1])
                     mnemonic += " " + candidates[0]
         mnemonic = mnemonic.strip()
         # check if mnemonic is valid
         if self.checker is not None and mnemonic is not None:
             if self.checker(mnemonic):
-                self.kb.clear_btn_ctrl(29, lv.btnm.CTRL.INACTIVE)
+                self.kb.clear_btn_ctrl(self.BTN_DONE, lv.btnm.CTRL.INACTIVE)
+            else:
+                self.kb.set_btn_ctrl(self.BTN_DONE, lv.btnm.CTRL.INACTIVE)
+            # check if we are at 12, 18 or 24 words
+            # offer to fix mnemonic if it's invalid
+            num_words = len(mnemonic.split())
+            if (
+                self.fixer is not None
+                and num_words in [12, 18, 24]
+                and self.kb.get_btn_ctrl(self.BTN_DONE, lv.btnm.CTRL.INACTIVE)
+            ):
+                # set correct button coordinates
+                y = -33 - self.table.get_height() // 2 if num_words == 18 else -38
+                x = -40 if num_words == 12 else -40 + self.table.get_width() // 2
+                # check if we can fix the mnemonic
+                try:
+                    self.fixer(mnemonic)
+                    self.fix_button.align(self.table, lv.ALIGN.OUT_BOTTOM_MID, x, y)
+                except:
+                    self.fix_button.align(
+                        self.table, lv.ALIGN.OUT_BOTTOM_MID, -400, -38
+                    )
             else:
-                self.kb.set_btn_ctrl(29, lv.btnm.CTRL.INACTIVE)
+                self.fix_button.align(self.table, lv.ALIGN.OUT_BOTTOM_MID, -400, -38)
 
     def callback(self, obj, event):
         if event != lv.EVENT.RELEASED:

src/gui/specter.py

@@ -30,14 +30,14 @@ async def new_mnemonic(self, generator):
         await self.load_screen(scr)
         return await scr.result()
 
-    async def recover(self, checker=None, lookup=None):
+    async def recover(self, checker=None, lookup=None, fix=None):
         """
         Asks the user for his recovery phrase.
         checker(mnemonic) - a function that validates recovery phrase
         lookup(word, num_candidates) - a function that
                 returns num_candidates words starting with word
         """
-        scr = RecoverMnemonicScreen(checker, lookup)
+        scr = RecoverMnemonicScreen(checker, lookup, fix)
         await self.load_screen(scr)
         return await scr.result()
 

src/keystore/memorycard.py

@@ -29,10 +29,7 @@ class MemoryCard(RAMKeyStore):
 In this mode device can only operate when the smartcard is inserted!"""
     # constants for secret storage
     MAGIC = b"sdiy\x00"  # specter-DIY version 0
-    KEYS = {
-        b"\x01": "enc",
-        b"\x02": "entropy",
-    }
+    KEYS = {b"\x01": "enc", b"\x02": "entropy"}
     # Button to go to storage menu
     # Menu should be implemented in async storage_menu function
     # Here we only have a single option - to show mnemonic

src/specter.py

@@ -181,15 +181,9 @@ async def initmenu(self):
         ]
         if self.keystore.is_key_saved and self.keystore.load_button:
             buttons.append((2, self.keystore.load_button))
-        buttons += [
-            (None, "Settings"),
-            (3, "Developer & USB settings"),
-        ]
+        buttons += [(None, "Settings"), (3, "Developer & USB settings")]
         if hasattr(self.keystore, "lock"):
-            buttons += [
-                (4, "Change PIN code"),
-                (5, "Lock device"),
-            ]
+            buttons += [(4, "Change PIN code"), (5, "Lock device")]
         # wait for menu selection
         menuitem = await self.gui.menu(buttons)
 
@@ -204,8 +198,13 @@ async def initmenu(self):
                 return self.mainmenu
         # recover
         elif menuitem == 1:
+            # a small function that fixes checksum of invalid mnemonic
+            def fix_mnemonic(phrase):
+                entropy = bip39.mnemonic_to_bytes(phrase, ignore_checksum=True)
+                return bip39.mnemonic_from_bytes(entropy)
+
             mnemonic = await self.gui.recover(
-                bip39.mnemonic_is_valid, bip39.find_candidates
+                bip39.mnemonic_is_valid, bip39.find_candidates, fix_mnemonic
             )
             if mnemonic is not None:
                 # load keys using mnemonic and empty password
@@ -253,24 +252,16 @@ async def mainmenu(self):
         buttons = (
             [
                 # id, text
-                (None, "Applications"),
+                (None, "Applications")
             ]
             + app_buttons
-            + [
-                (None, "Communication"),
-            ]
+            + [(None, "Communication")]
             + host_buttons
-            + [
-                (None, "More"),  # delimiter
-            ]
+            + [(None, "More")]  # delimiter
         )
         if hasattr(self.keystore, "lock"):
-            buttons += [
-                (2, "Lock device"),
-            ]
-        buttons += [
-            (3, "Settings"),
-        ]
+            buttons += [(2, "Lock device")]
+        buttons += [(3, "Settings")]
         # wait for menu selection
         menuitem = await self.gui.menu(buttons)
 
@@ -311,23 +302,10 @@ async def settingsmenu(self):
         ]
         if self.keystore.storage_button is not None:
             buttons.append((0, self.keystore.storage_button))
-        buttons.extend(
-            [
-                (2, "Enter password"),
-                (None, "Security"),  # delimiter
-            ]
-        )
+        buttons.extend([(2, "Enter password"), (None, "Security")])  # delimiter
         if hasattr(self.keystore, "lock"):
-            buttons.extend(
-                [
-                    (3, "Change PIN code"),
-                ]
-            )
-        buttons.extend(
-            [
-                (4, "Developer & USB"),
-            ]
-        )
+            buttons.extend([(3, "Change PIN code")])
+        buttons.extend([(4, "Developer & USB")])
         # wait for menu selection
         menuitem = await self.gui.menu(buttons, last=(255, None))
 
@@ -463,10 +441,7 @@ def load_config(self):
                 print(e)
 
     def update_config(self, usb=False, dev=False, **kwargs):
-        config = {
-            "usb": usb,
-            "dev": dev,
-        }
+        config = {"usb": usb, "dev": dev}
         self.keystore.save_aead(
             self.path + "/settings",
             adata=json.dumps(config).encode(),