Add feature to export multiple accounts XPub keys (#258)

* Add feature to export multiple accounts XPub keys

* Apply suggestions from code review

still need to do a few modifs and test them.

Co-authored-by: Stepan Snigirev <[email protected]>

* refactor pr to avoid code duplication

---------

Co-authored-by: polto <[email protected]>
Co-authored-by: Stepan Snigirev <[email protected]>

Author: 3 people

src/apps/xpubs/xpubs.py

@@ -5,10 +5,9 @@
 from binascii import hexlify
 from embit.liquid.networks import NETWORKS
 from embit import bip32
-from helpers import is_liquid
+from helpers import is_liquid, SDCardFile
 from io import BytesIO
 import platform
-import lvgl as lv
 from collections import OrderedDict
 
 class XpubApp(BaseApp):
@@ -41,7 +40,8 @@ async def menu(self, show_screen, show_all=False):
                 (0, "Show more keys"),
                 (2, "Change account number"),
                 (1, "Enter custom derivation"),
-                (3, "Export all keys to SD"),
+                (3, "Export all keys for this account"),
+                (4, "Export multiple accounts"),
             ]
         else:
             buttons = [
@@ -69,9 +69,14 @@ async def menu(self, show_screen, show_all=False):
                 ),
             ]
         # wait for menu selection
-        menuitem = await show_screen(Menu(buttons, last=(255, None),
-                                          title="Select the key",
-                                          note="Current account number: %d" % self.account))
+        menuitem = await show_screen(
+            Menu(
+                buttons,
+                title="Select the key",
+                note="Current account number: %d" % self.account,
+                last=(255, None),
+            )
+        )
 
         # process the menu button:
         # back button
@@ -84,6 +89,15 @@ async def menu(self, show_screen, show_all=False):
             if der is not None:
                 await self.show_xpub(der, show_screen)
                 return True
+        elif menuitem == 2:
+            account = await show_screen(NumericScreen(current_val=str(self.account)))
+            if account and int(account) >= 0x80000000:
+                raise AppError("Account number too large")
+            try:
+                self.account = int(account)
+            except:
+                self.account = 0
+            return await self.menu(show_screen)
         elif menuitem == 3:
             file_format = await self.save_menu(show_screen)
             if file_format:
@@ -93,66 +107,143 @@ async def menu(self, show_screen, show_all=False):
                           "Public keys are saved to the file:\n\n%s" % filename,
                           button_text="Close")
                 )
-        elif menuitem == 2:
-            account = await show_screen(NumericScreen(current_val=str(self.account)))
-            if account and int(account) > 0x80000000:
-                raise AppError('Account number too large')
-            try:
-                self.account = int(account)
-            except:
-                self.account = 0
-            return await self.menu(show_screen)
+        elif menuitem == 4:
+            from_account = await show_screen(
+                NumericScreen(
+                    title="Enter START account number",
+                    current_val=str(self.account)
+                )
+            )
+            to_account = await show_screen(
+                NumericScreen(
+                    title="Enter END account number",
+                    current_val=str(self.account)
+                )
+            )
+            if from_account is None or to_account is None:
+                return
+            if from_account == "":
+                from_account = self.account
+            if to_account == "":
+                to_account = self.account
+            from_account = int(from_account)
+            to_account = int(to_account)
+            file_format = await self.save_menu(show_screen)
+            await self.export_multiple_accounts_xpubs(
+                from_account, to_account, file_format, show_screen
+            )
         else:
             await self.show_xpub(menuitem, show_screen)
             return True
         return False
 
-    def save_all_to_sd(self, file_format):
+    def save_all_to_sd(self, file_format, account=None):
+        if account is None:
+            account = self.account
+
         fingerprint = hexlify(self.keystore.fingerprint).decode()
-        extension = "json"
-        coin = NETWORKS[self.network]["bip32"]
 
+        extension = "txt" if file_format == self.export_specter_diy else "json"
+        filename = "%s-%s-%d-all.%s" % (
+            file_format, fingerprint, account, extension,
+        )
+
+        if not platform.is_sd_present():
+            raise AppError("Please insert SD card")
+
+        with SDCardFile(filename, "w") as f:
+            self._dump_account(f, file_format, account)
+
+        return filename
+
+    async def export_multiple_accounts_xpubs(
+        self,
+        from_account,
+        to_account,
+        file_format,
+        show_screen,
+    ):
+        if from_account > to_account:
+            from_account, to_account = to_account, from_account
+        if to_account >= 0x80000000:
+            raise AppError('Account number too large')
+        fingerprint = hexlify(self.keystore.fingerprint).decode()
+        if file_format == self.export_specter_diy:
+            # in our format we can dump any number of accounts in one file
+            filename = "%s-%s-%d-%d.txt" % (
+                file_format, fingerprint, from_account, to_account
+            )
+            with SDCardFile(filename, "w") as f:
+                for account in range(from_account, to_account+1):
+                    self.show_loader(title="Exporting account %d..." % account)
+                    self._dump_account(f, file_format, account)
+            await show_screen(
+                Alert(
+                    "Success!",
+                    "File was successfully saved under:\n\n%s" % filename,
+                    button_text="OK",
+                )
+            )
+        else: # cc format - one file per account
+            for account in range(from_account, to_account+1):
+                self.show_loader(title="Exporting account %d..." % account)
+                self.save_all_to_sd(file_format, account)
+            await show_screen(
+                Alert(
+                    "Success!",
+                    "All accounts are saved to corresponding files.",
+                    button_text="OK",
+                )
+            )
+
+    def _dump_account(self, f, file_format, account):
+        """dump all keys of one account to a file"""
+        coin = NETWORKS[self.network]["bip32"]
         derivations = [
-            ('bip84', "p2wpkh", "m/84'/%d'/%d'" % (coin, self.account)),
-            ('bip86', "p2tr", "m/86'/%d'/%d'" % (coin, self.account)),
-            ('bip49', "p2sh-p2wpkh", "m/49'/%d'/%d'" % (coin, self.account)),
-            ('bip44', "p2pkh", "m/44'/%d'/%d'" % (coin, self.account)),
-            ('bip48_1', "p2sh-p2wsh", "m/48'/%d'/%d'/1'" % (coin, self.account)),
-            ('bip48_2', "p2wsh", "m/48'/%d'/%d'/2'" % (coin, self.account)),
+            ('bip84', "p2wpkh", "m/84'/%d'/%d'" % (coin, account)),
+            ('bip86', "p2tr", "m/86'/%d'/%d'" % (coin, account)),
+            ('bip49', "p2sh-p2wpkh", "m/49'/%d'/%d'" % (coin, account)),
+            ('bip44', "p2pkh", "m/44'/%d'/%d'" % (coin, account)),
+            ('bip48_1', "p2sh-p2wsh", "m/48'/%d'/%d'/1'" % (coin, account)),
+            ('bip48_2', "p2wsh", "m/48'/%d'/%d'/2'" % (coin, account)),
         ]
+        fingerprint = hexlify(self.keystore.fingerprint).decode()
 
         if file_format == self.export_specter_diy:
-            # text file with [fgp/der]xpub lines
-            filedata = ""
-            for der in derivations:
-                xpub = self.keystore.get_xpub(der[2])
-                filedata += "[%s/%s]%s\n" % (fingerprint, der[2][2:].replace("'","h"), xpub.to_base58(NETWORKS[self.network]["xpub"]))
-            extension = "txt"
+            for keytype, scripttype, der in derivations:
+                xpub = self.keystore.get_xpub(der)
+                f.write("[%s/%s]%s\n" % (
+                    fingerprint,
+                    der.replace("m/", "").replace("'","h"),
+                    xpub.to_base58(NETWORKS[self.network]["xpub"]),
+                ))
         else:
             # coldcard generic json format
             m = self.keystore.get_xpub("m")
             data = {
                 "xpub": m.to_base58(NETWORKS[self.network]["xpub"]),
                 "xfp": fingerprint,
-                "account": self.account,
+                "account": account,
                 "chain": "BTC" if self.network == "main" else "XTN"
             }
 
-            for der in derivations:
-                xpub = self.keystore.get_xpub(der[2])
+            for keytype, scripttype, der in derivations:
+                xpub = self.keystore.get_xpub(der)
 
-                data[der[0]] = {
-                    "name": der[1],
-                    "deriv": der[2],
+                data[keytype] = {
+                    "name": scripttype,
+                    "deriv": der,
                     "xpub": xpub.to_base58(NETWORKS[self.network]["xpub"]),
-                    "_pub": xpub.to_base58(bip32.detect_version(der[2], default="xpub", network=NETWORKS[self.network]))
+                    "_pub": xpub.to_base58(
+                        bip32.detect_version(
+                            der,
+                            default="xpub",
+                            network=NETWORKS[self.network]
+                        )
+                    )
                 }
 
-            filedata = json.dumps(data).encode()
-
-        filename = "%s-%s-%d-all.%s" % (file_format, fingerprint, self.account, extension)
-        self.write_file(filename, filedata)
-        return filename
+            json.dump(data, f)
 
     async def process_host_command(self, stream, show_screen):
         if self.keystore.is_locked:
@@ -192,12 +283,15 @@ async def show_xpub(self, derivation, show_screen):
             fingerprint,
             derivation[1:],
         )
-        res = await show_screen(XPubScreen(xpub=canonical, slip132=slip132, prefix=prefix))
+        res = await show_screen(
+            XPubScreen(xpub=canonical, slip132=slip132, prefix=prefix)
+        )
         if res == XPubScreen.CREATE_WALLET:
             await self.create_wallet(derivation, canonical, prefix, ver, show_screen)
         elif res:
             filename = "%s-%s.txt" % (fingerprint, derivation[2:].replace("/", "-"))
-            self.write_file(filename, res)
+            with SDCardFile(filename, "w") as f:
+                f.write(res)
             await show_screen(
                 Alert("Saved!",
                       "Extended public key is saved to the file:\n\n%s" % filename,
@@ -283,14 +377,6 @@ async def create_wallet(self, derivation, xpub, prefix, version, show_screen):
             await self.communicate(stream, app="wallets")
 
 
-    def write_file(self, filename, filedata):
-        if not platform.is_sd_present():
-            raise AppError("SD card is not present")
-        platform.mount_sdcard()
-        with open(platform.fpath("/sd/%s" % filename), "w") as f:
-            f.write(filedata)
-        platform.unmount_sdcard()
-
     async def save_menu(self, show_screen):
         buttons = [(0, "Specter-DIY (plaintext)"), (1, "Cold Card (json)")]
         # wait for menu selection
@@ -307,6 +393,7 @@ async def save_menu(self, show_screen):
             return self.export_coldcard
         return None
 
+
     def wipe(self):
         # nothing to delete
         pass

src/helpers.py

@@ -174,3 +174,38 @@ def read_write(fin, fout, chunk_size=32):
         chunk = fin.read(chunk_size)
         total += fout.write(chunk)
     return total
+
+class SDCardFile:
+    """open file on SD card, mount card on enter, unmount on exit
+
+    Usage:
+
+    with SDCardFile("/sd/blah", "w") as f:
+        f.write("blah")
+    """
+    def __init__(self, filename, *args, **kwargs):
+        self._filename = filename
+        self._args = args
+        self._kwargs = kwargs
+        self.file = None
+
+    def mount(self):
+        platform.mount_sdcard()
+
+    def unmount(self):
+        platform.unmount_sdcard()
+
+    def __enter__(self):
+        self.mount()
+        self.file = open(
+            platform.fpath("/sd/" + self._filename),
+            *self._args, **self._kwargs
+        )
+        return self.file
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.file:
+            self.file.close()
+            self.file = None
+        self.unmount()
+

src/platform.py

@@ -288,4 +288,4 @@ def get_battery_status():
         return level, charging
     except Exception as e:
         print(e)
-        return None, None
+        return None, None