Encryption load error now shows only one error message

selfcustody · Nov 21, 2023 · 4f03575 · 4f03575
1 parent 4857f4b
commit 4f03575
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 13 deletions.
diff --git a/src/krux/pages/encryption_ui.py b/src/krux/pages/encryption_ui.py
@@ -55,7 +55,7 @@ def encryption_key(self):
         )
         _, key = submenu.run_loop()
         if key in (ESC_KEY, MENU_CONTINUE):
-            return None
+            return ""
 
         if key:
             self.ctx.display.clear()
@@ -277,14 +277,11 @@ def _load_encrypted_mnemonic(self, mnemonic_id, sd_card=False):
 
         key_capture = EncryptionKey(self.ctx)
         key = key_capture.encryption_key()
-        if key is None:
-            self.flash_text(t("Failed to decrypt"), theme.error_color)
+        if key in (None, "", ESC_KEY):
+            self.flash_text(t("Key was not provided"), theme.error_color)
             return MENU_CONTINUE
         self.ctx.display.clear()
         self.ctx.display.draw_centered_text(t("Processing ..."))
-        if key in ("", ESC_KEY):
-            self.flash_text(t("Failed to decrypt"), theme.error_color)
-            return MENU_CONTINUE
         mnemonic_storage = MnemonicStorage()
         try:
             words = mnemonic_storage.decrypt(key, mnemonic_id, sd_card).split()

diff --git a/src/krux/pages/login.py b/src/krux/pages/login.py
@@ -430,19 +430,17 @@ def _encrypted_qr_code(self, data):
 
                 key_capture = EncryptionKey(self.ctx)
                 key = key_capture.encryption_key()
-                if key is None:
-                    self.flash_text(t("Mnemonic was not decrypted"))
-                    return None
+                if key in (None, "", ESC_KEY):
+                    self.flash_text(t("Key was not provided"), theme.error_color)
+                    return MENU_CONTINUE
                 self.ctx.display.clear()
                 self.ctx.display.draw_centered_text(t("Processing ..."))
-                if key in ("", ESC_KEY):
-                    self.flash_text(t("Failed to decrypt"), theme.error_color)
-                    return None
                 word_bytes = encrypted_qr.decrypt(key)
                 if word_bytes is None:
                     self.flash_text(t("Failed to decrypt"), theme.error_color)
-                    return None
+                    return MENU_CONTINUE
                 return bip39.mnemonic_from_bytes(word_bytes).split()
+            return MENU_CONTINUE  # prompt NO
         return None
 
     def load_key_from_qr_code(self):
@@ -490,6 +488,8 @@ def load_key_from_qr_code(self):
                     ]
             if not words:
                 words = self._encrypted_qr_code(data)
+                if words == MENU_CONTINUE:
+                    return MENU_CONTINUE
         if not words or (len(words) != 12 and len(words) != 24):
             self.flash_text(t("Invalid mnemonic length"), theme.error_color)
             return MENU_CONTINUE

diff --git a/tests/pages/test_encryption_ui.py b/tests/pages/test_encryption_ui.py
@@ -72,6 +72,7 @@ def test_load_key_from_qr_code(m5stickv, mocker):
     from krux.pages.encryption_ui import EncryptionKey, ENCRYPTION_KEY_MAX_LEN
     from krux.input import BUTTON_ENTER, BUTTON_PAGE
 
+    print("case 1: load_key_from_qr_code")
     BTN_SEQUENCE = (
         [BUTTON_PAGE]  # move to QR code key
         + [BUTTON_ENTER]  # choose QR code key
@@ -87,6 +88,7 @@ def test_load_key_from_qr_code(m5stickv, mocker):
     key = key_generator.encryption_key()
     assert key == "qr key"
 
+    print("case 2: load_key_from_qr_code")
     # Repeat with too much characters >ENCRYPTION_KEY_MAX_LEN
     BTN_SEQUENCE = [BUTTON_PAGE] + [  # move to QR code key
         BUTTON_ENTER