Merge pull request #15 from odudex/new_features

Store/load/delete mnemonic on both flash and SD

pyproject.toml

@@ -22,7 +22,7 @@
 
 [tool.poetry]
 name = "krux"
-version = "23.04.beta1"
+version = "23.04.beta3"
 description = "Open-source signing device firmware for Bitcoin"
 authors = ["Jeff S <[email protected]>"]
 

src/krux/encryption.py

@@ -29,19 +29,24 @@
 from .baseconv import base_encode, base_decode
 from .sd_card import SDHandler
 
-SEED_FILE = "seeds.json"
-
+MNEMONICS_FILE = "seeds.json"
+MODE_ECB = 3
 
 class AESCipher:
     """Helper for AES encrypt/decrypt"""
 
-    def __init__(self, key):
-        self.key = hashlib.sha256(key.encode()).digest()
+    def __init__(self, key, salt):
+        self.key = hashlib.pbkdf2_hmac(
+            'sha256',
+            key.encode(),
+            salt.encode(),
+            100000
+        )
 
     def encrypt(self, raw):
         """Encrypt using AES MODE_ECB and return the value encoded as base64"""
         data_bytes = raw.encode()
-        encryptor = ucryptolib.aes(self.key, ucryptolib.MODE_ECB)
+        encryptor = ucryptolib.aes(self.key, MODE_ECB)
         encrypted = encryptor.encrypt(
             data_bytes + b"\x00" * ((16 - (len(data_bytes) % 16)) % 16)
         )
@@ -50,65 +55,103 @@ def encrypt(self, raw):
     def decrypt(self, enc):
         """Decrypt a base64 using AES MODE_ECB and return the value decoded as string"""
         encrypted = base_decode(enc, 64)  # test - decode 64
-        decryptor = ucryptolib.aes(self.key, ucryptolib.MODE_ECB)
+        decryptor = ucryptolib.aes(self.key, MODE_ECB)
         load = decryptor.decrypt(encrypted).decode("utf-8")
         return load.replace("\x00", "")
 
 
-class StoredSeeds:
+class MnemonicStorage:
     """Handler of stored encrypted seeds"""
 
     def __init__(self) -> None:
-        self.encrypted_store = {}
+        self.stored = {}
+        self.stored_sd = {}
+        self.has_sd_card = False
         try:
             with SDHandler() as sd:
-                self.encrypted_store = json.loads(sd.read(SEED_FILE))
+                self.stored_sd = json.loads(sd.read(MNEMONICS_FILE))
+                self.has_sd_card = True
+        except:
+            pass
+        try:
+            with open("/flash/" + MNEMONICS_FILE, "r") as f:
+                self.stored = json.loads(f.read())
         except:
             pass
 
-    def list_fingerprints(self):
+
+    def list_mnemonics(self, sd_card=False):
         """List all seeds stored on a file"""
-        fingerprints = []
-        for fingerprint in self.encrypted_store:
-            fingerprints.append(fingerprint)
-        return fingerprints
-
-    def decrypt(self, key, fingerprint):
-        """Decrypt a selected seed from a file"""
-        decryptor = AESCipher(key)
+        mnemonic_ids = []
+        source = self.stored_sd if sd_card else self.stored
+        for mnemonic_id in source:
+            mnemonic_ids.append(mnemonic_id)
+        return mnemonic_ids
+
+    def decrypt(self, key, mnemonic_id, sd_card=False):
+        """Decrypt a selected encrypted mnemonic from a file"""
+        decryptor = AESCipher(key, mnemonic_id)
         try:
-            load = self.encrypted_store.get(fingerprint)
+            if sd_card:
+                load = self.stored_sd.get(mnemonic_id)
+            else:
+                load = self.stored.get(mnemonic_id)
             words = decryptor.decrypt(load)
         except:
             return None
         return words
 
-    def store_encrypted(self, key, fingerprint, seed):
-        """Saves the seed encrypted on a file"""
-        encryptor = AESCipher(key)
-        encrypted = encryptor.encrypt(seed).decode("utf-8")
-        seeds = {}
-
-        # load current SEED_FILE
-        try:
-            with SDHandler() as sd:
-                seeds = json.loads(sd.read(SEED_FILE))
-        except:
-            pass
-
-        # save the new SEED_FILE
-        try:
-            with SDHandler() as sd:
-                seeds[fingerprint] = encrypted
-                sd.write(SEED_FILE, json.dumps(seeds))
-        except:
-            pass
+    def store_encrypted(self, key, mnemonic_id, mnemonic, sd_card=False):
+        """Saves the encrypted mnemonic on a file"""
+        encryptor = AESCipher(key, mnemonic_id)
+        encrypted = encryptor.encrypt(mnemonic).decode("utf-8")
+        mnemonics = {}
+        success = True
+        if sd_card:
+            # load current MNEMONICS_FILE
+            try:
+                with SDHandler() as sd:
+                    mnemonics = json.loads(sd.read(MNEMONICS_FILE))
+            except:
+                success = False
+
+            # save the new MNEMONICS_FILE
+            try:
+                with SDHandler() as sd:
+                    mnemonics[mnemonic_id] = encrypted
+                    sd.write(MNEMONICS_FILE, json.dumps(mnemonics))
+            except:
+                success = False
+        else:
+            try:
+                # save the new SETTINGS_FILENAME
+                with open("/flash/" + MNEMONICS_FILE, "r") as f:
+                    mnemonics = json.loads(f.read())
+            except:
+                success = False
+            try:
+                with open("/flash/" + MNEMONICS_FILE, "w") as f:
+                    mnemonics[mnemonic_id] = encrypted
+                    f.write(json.dumps(mnemonics))
+            except:
+                success = False
+        return success
+
+
+    def del_mnemonic(self, mnemonic_id, sd_card=False):
+        """Remove an entry from encrypted mnemonics file"""
+        if sd_card:
+            self.stored_sd.pop(mnemonic_id)
+            try:
+                with SDHandler() as sd:
+                    sd.write(MNEMONICS_FILE, json.dumps(self.stored_sd))
+            except:
+                pass
+        else:
+            self.stored.pop(mnemonic_id)
+            try:
+                with open("/flash/" + MNEMONICS_FILE, "w") as f:
+                    f.write(json.dumps(self.stored))
+            except:
+                pass
 
-    def del_seed(self, fingerprint):
-        """Remove an entry from encrypted seeds file"""
-        self.encrypted_store.pop(fingerprint)
-        try:
-            with SDHandler() as sd:
-                sd.write(SEED_FILE, json.dumps(self.encrypted_store))
-        except:
-            pass

src/krux/metadata.py

@@ -19,5 +19,5 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
-VERSION = "23.04.beta1"
+VERSION = "23.04.beta3"
 SIGNER_PUBKEY = "03339e883157e45891e61ca9df4cd3bb895ef32d475b8e793559ea10a36766689b"

src/krux/pages/home.py

@@ -57,7 +57,7 @@
 import qrcode
 from ..printers import create_printer
 from ..printers.cnc import FilePrinter
-from ..encryption import StoredSeeds
+from ..encryption import MnemonicStorage
 import board
 import uos
 import time
@@ -98,6 +98,7 @@ def __init__(self, ctx):
                 ctx,
                 [
                     (t("Mnemonic"), self.mnemonic),
+                    (t("Encrypt Mnemonic"), self.encrypt_mnemonic),
                     (t("Extended Public Key"), self.public_key),
                     (t("Wallet Descriptor"), self.wallet),
                     (t("Address"), self.list_address),
@@ -125,7 +126,6 @@ def mnemonic(self):
                 (t("SeedQR"), self.display_seed_qr),
                 (t("Stackbit 1248"), self.stackbit),
                 (t("Tiny Seed"), self.tiny_seed),
-                (t("Store Encrypted Seed"), self.store_encrypted_seed),
                 (t("Back"), lambda: MENU_EXIT),
             ],
         )
@@ -420,33 +420,59 @@ def tiny_seed(self):
                 tiny_seed.print_tiny_seed()
         return MENU_CONTINUE
 
-    def store_encrypted_seed(self):
-        """Handler for Mnemonic > Store Encrypted Seed menu item"""
-        fingerprint = self.ctx.wallet.key.fingerprint_hex_str()
-        stored_seeds = StoredSeeds()
-        if fingerprint in stored_seeds.list_fingerprints():
-            if self.prompt(
-                t("Seed already stored, would you like to delete it?"),
+    def store_mnemonic_on_memory(self, sd_card=False):
+        key = self.capture_from_keypad(
+            t("Encryption Key"),
+            [LETTERS, UPPERCASE_LETTERS, NUM_SPECIAL_1, NUM_SPECIAL_2],
+        )
+        if key in ("", ESC_KEY):
+            self.ctx.display.flash_text(t("Encrypted mnemonic was not stored"))
+            return
+        self.ctx.display.clear()
+        if self.prompt(
+                t("Give this mnemonic a custom ID? Otherwise current fingerprint will be used"),
                 self.ctx.display.height() // 2,
             ):
-                stored_seeds.del_seed(fingerprint)
+            mnemonic_id = self.capture_from_keypad(
+                t("Mnemonic Storage ID"),
+                [LETTERS, UPPERCASE_LETTERS, FILE_SPECIAL],
+            )
         else:
-            key = self.capture_from_keypad(
-                t("Encryption Key"),
-                [LETTERS, UPPERCASE_LETTERS, NUM_SPECIAL_1, NUM_SPECIAL_2],
+            mnemonic_id = self.ctx.wallet.key.fingerprint_hex_str()
+        words = self.ctx.wallet.key.mnemonic
+        mnemonic_storage = MnemonicStorage()
+        self.ctx.display.clear()
+        self.ctx.display.draw_centered_text( t("Processing ..."))
+        if mnemonic_storage.store_encrypted(key, mnemonic_id, words, sd_card):
+            self.ctx.display.clear()
+            self.ctx.display.draw_centered_text(
+                t("Encrypted mnemonic was stored with ID: %s" % mnemonic_id)
             )
-            if key in ("", ESC_KEY):
-                self.ctx.display.flash_text(t("Encrypted seed was not stored"))
-                del stored_seeds
-                return  # MENU_CONTINUE
-            words = self.ctx.wallet.key.mnemonic
-            stored_seeds.store_encrypted(key, fingerprint, words)
+        else:
             self.ctx.display.clear()
             self.ctx.display.draw_centered_text(
-                t("Encrypted seed was stored with ID: %s" % fingerprint)
+                t("Failed to store mnemonic")
             )
-            self.ctx.input.wait_for_button()
-        del stored_seeds
+        self.ctx.input.wait_for_button()
+        del mnemonic_storage
+
+    def encrypt_mnemonic(self):
+        """Handler for Mnemonic > Encrypt Mnemonic menu item"""
+        encrypt_outputs_menu = []
+        encrypt_outputs_menu.append(
+            (t("Store on Flash"), self.store_mnemonic_on_memory)
+        )
+        mnemonic_storage = MnemonicStorage()
+        if mnemonic_storage.has_sd_card:
+            encrypt_outputs_menu.append(
+            (t("Store on SD Card"), lambda: self.store_mnemonic_on_memory(sd_card=True))
+        )
+        del mnemonic_storage
+        encrypt_outputs_menu.append((t("Back"), lambda: MENU_EXIT))
+        submenu = Menu(self.ctx, encrypt_outputs_menu)
+        _, _ = submenu.run_loop()
+        return MENU_CONTINUE
+
 
     def public_key(self):
         """Handler for the 'xpub' menu item"""