From 4f0357508cd9419ac380ed64d112c389d2677d31 Mon Sep 17 00:00:00 2001
From: tadeubas <tadeubas@gmail.com>
Date: Tue, 21 Nov 2023 15:06:03 -0300
Subject: [PATCH] Encryption load error now shows only one error message

---
 src/krux/pages/encryption_ui.py   |  9 +++------
 src/krux/pages/login.py           | 14 +++++++-------
 tests/pages/test_encryption_ui.py |  2 ++
 3 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/src/krux/pages/encryption_ui.py b/src/krux/pages/encryption_ui.py
index ea224e78a..976eff831 100644
--- 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
index b03383260..ee90f37ce 100644
--- 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
index ffeb8d87f..016eb1b35 100644
--- 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