Merge pull request #290 from tadeubas/integrated_changes

Encryption load error now shows only one error message

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()

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

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

tests/test_encryption.py

@@ -61,6 +61,9 @@ def mock_file_operations(mocker):
     mocker.patch("builtins.open", mocker.mock_open(read_data=SEEDS_JSON))
 
 
+# -------------------------
+
+
 def test_ecb_encryption(m5stickv):
     from krux.encryption import AESCipher
 
@@ -117,7 +120,7 @@ def test_load_decrypt_cbc(m5stickv, mock_file_operations):
     assert words_sd == CBC_WORDS
 
 
-def test_encrypt_ecb_flash(mocker):
+def test_encrypt_ecb_flash(m5stickv, mocker):
     from krux.krux_settings import Settings
     from krux.encryption import MnemonicStorage
 
@@ -129,7 +132,7 @@ def test_encrypt_ecb_flash(mocker):
     m().write.assert_called_once_with(ECB_ONLY_JSON)
 
 
-def test_encrypt_cbc_flash(mocker):
+def test_encrypt_cbc_flash(m5stickv, mocker):
     from krux.krux_settings import Settings
     from krux.encryption import MnemonicStorage
 
@@ -143,7 +146,7 @@ def test_encrypt_cbc_flash(mocker):
     m().write.assert_called_once_with(CBC_ONLY_JSON)
 
 
-def test_encrypt_ecb_sd(mocker, mock_file_operations):
+def test_encrypt_ecb_sd(m5stickv, mocker, mock_file_operations):
     from krux.krux_settings import Settings
     from krux.encryption import MnemonicStorage
 
@@ -155,7 +158,7 @@ def test_encrypt_ecb_sd(mocker, mock_file_operations):
     m().write.assert_called_once_with(ECB_ONLY_JSON)
 
 
-def test_encrypt_cbc_sd(mocker, mock_file_operations):
+def test_encrypt_cbc_sd(m5stickv, mocker, mock_file_operations):
     from krux.krux_settings import Settings
     from krux.encryption import MnemonicStorage
 
@@ -169,7 +172,7 @@ def test_encrypt_cbc_sd(mocker, mock_file_operations):
     m().write.assert_called_once_with(CBC_ONLY_JSON)
 
 
-def test_delet_from_flash(mocker):
+def test_delet_from_flash(m5stickv, mocker):
     from krux.encryption import MnemonicStorage
 
     # Loads a file with 2 mnemonics, one with ID="ecbID", other with ID="cbcID"
@@ -180,7 +183,7 @@ def test_delet_from_flash(mocker):
     m().write.assert_called_once_with(CBC_ONLY_JSON)
 
 
-def test_delet_from_sd(mocker, mock_file_operations):
+def test_delet_from_sd(m5stickv, mocker, mock_file_operations):
     from krux.encryption import MnemonicStorage
 
     # Loads a file with 2 mnemonics, one with ID="ecbID", other with ID="cbcID"
@@ -191,7 +194,7 @@ def test_delet_from_sd(mocker, mock_file_operations):
     m().write.assert_called_once_with(CBC_ONLY_JSON)
 
 
-def test_create_ecb_encrypted_qr_code():
+def test_create_ecb_encrypted_qr_code(m5stickv):
     from krux.encryption import EncryptedQRCode
     from krux.krux_settings import Settings
 
@@ -201,7 +204,7 @@ def test_create_ecb_encrypted_qr_code():
     assert qr_data == ECB_ENCRYPTED_QR
 
 
-def test_create_cbc_encrypted_qr_code():
+def test_create_cbc_encrypted_qr_code(m5stickv):
     from krux.encryption import EncryptedQRCode
     from krux.krux_settings import Settings
 
@@ -212,7 +215,7 @@ def test_create_cbc_encrypted_qr_code():
     assert qr_data == CBC_ENCRYPTED_QR
 
 
-def test_decode_ecb_encrypted_qr_code():
+def test_decode_ecb_encrypted_qr_code(m5stickv):
     from krux.encryption import EncryptedQRCode
     from embit import bip39
 
@@ -224,7 +227,7 @@ def test_decode_ecb_encrypted_qr_code():
     assert words == TEST_WORDS
 
 
-def test_decode_cbc_encrypted_qr_code():
+def test_decode_cbc_encrypted_qr_code(m5stickv):
     from krux.encryption import EncryptedQRCode
     from embit import bip39
 
@@ -237,14 +240,24 @@ def test_decode_cbc_encrypted_qr_code():
     assert words == TEST_WORDS
 
 
-def create_ctx(mocker, btn_seq, touch_seq=None):
-    """Helper to create mocked context obj"""
-    ctx = mock_context(mocker)
-    ctx.power_manager.battery_charge_remaining.return_value = 1
-    ctx.input.wait_for_button = mocker.MagicMock(side_effect=btn_seq)
+def test_customize_pbkdf2_iterations_create_and_decode(m5stickv):
+    from krux.encryption import EncryptedQRCode
+    from krux.krux_settings import Settings
+    from embit import bip39
 
-    if touch_seq:
-        ctx.input.touch = mocker.MagicMock(
-            current_index=mocker.MagicMock(side_effect=touch_seq)
-        )
-    return ctx
+    print("case Encode: customize_pbkdf2_iterations")
+    Settings().encryption.version = "AES-ECB"
+    Settings().encryption.pbkdf2_iterations = 99999
+    encrypted_qr = EncryptedQRCode()
+    qr_data = encrypted_qr.create(TEST_KEY, TEST_MNEMONIC_ID, TEST_WORDS)
+    print(qr_data)
+    print(ECB_ENCRYPTED_QR)
+
+    print("case Decode: customize_pbkdf2_iterations")
+    public_data = encrypted_qr.public_data(qr_data)
+    assert public_data == (
+        "Encrypted QR Code:\nID: test ID\nVersion: AES-ECB\nKey iter.: 90000"
+    )
+    word_bytes = encrypted_qr.decrypt(TEST_KEY)
+    words = bip39.mnemonic_from_bytes(word_bytes)
+    assert words == TEST_WORDS