taproot miniscript refactor

selfcustody · Dec 19, 2024 · db8981c · db8981c
1 parent d10b979
commit db8981c
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 23 deletions.
diff --git a/src/krux/pages/wallet_settings.py b/src/krux/pages/wallet_settings.py
@@ -257,7 +257,7 @@ def _miniscript_type(self):
             self.ctx,
             [
                 ("Native Segwit - P2WSH", lambda: P2WSH),
-                ("Taproot - TR", lambda: P2TR),
+                ("Taproot - P2TR", lambda: P2TR),
             ],
             disable_statusbar=True,
         )

diff --git a/src/krux/psbt.py b/src/krux/psbt.py
@@ -736,7 +736,10 @@ def get_policy(scope, scriptpubkey, xpubs):
             cosigners = get_cosigners_taproot_miniscript(
                 scope.taproot_bip32_derivations, xpubs
             )
-            policy.update({"cosigners": cosigners, "miniscript": P2TR})
+            policy.update({"cosigners": cosigners})
+            if len(cosigners) > 1:
+                # Assume it is single-sig TR if there is only one cosigner
+                policy.update({"miniscript": P2TR})
         except Exception as e:
             print("Error getting taproot PSBT cosigners: ", e)
 

diff --git a/src/krux/wallet.py b/src/krux/wallet.py
@@ -118,6 +118,30 @@ def is_loaded(self):
         """Returns a boolean indicating whether or not this wallet has been loaded"""
         return self.wallet_data is not None
 
+    def _validate_descriptor(self, descriptor, descriptor_xpubs):
+        """Validates the descriptor against the current key and policy type"""
+
+        if self.is_multisig():
+            if not descriptor.is_basic_multisig:
+                raise ValueError("not multisig")
+            if self.key.xpub() not in descriptor_xpubs:
+                raise ValueError("xpub not a multisig cosigner")
+        elif self.is_miniscript():
+            if self.key.script_type == P2WSH:
+                if descriptor.miniscript is None or descriptor.is_basic_multisig:
+                    raise ValueError("not P2WSH miniscript")
+            elif self.key.script_type == P2TR:
+                if descriptor.taptree is None:
+                    raise ValueError("not P2TR miniscript")
+            if self.key.xpub() not in descriptor_xpubs:
+                raise ValueError("xpub not a miniscript cosigner")
+        else:
+            if not descriptor.key:
+                if len(descriptor.keys) > 1:
+                    raise ValueError("not single-sig")
+            if self.key.xpub() != descriptor_xpubs[0]:
+                raise ValueError("xpub does not match")
+
     def load(self, wallet_data, qr_format, allow_assumption=None):
         """Loads the wallet from the given data"""
         descriptor, label = parse_wallet(wallet_data, allow_assumption)
@@ -131,26 +155,10 @@ def load(self, wallet_data, qr_format, allow_assumption=None):
             )
 
         if self.key:
-            if self.is_multisig():
-                if not descriptor.is_basic_multisig:
-                    raise ValueError("not multisig")
-                if self.key.xpub() not in descriptor_xpubs:
-                    raise ValueError("xpub not a multisig cosigner")
-            elif self.is_miniscript():
-                if self.key.script_type == P2WSH:
-                    if descriptor.miniscript is None or descriptor.is_basic_multisig:
-                        raise ValueError("not P2WSH miniscript")
-                elif self.key.script_type == P2TR:
-                    if descriptor.taptree is None:
-                        raise ValueError("not P2TR miniscript")
-                if self.key.xpub() not in descriptor_xpubs:
-                    raise ValueError("xpub not a miniscript cosigner")
-            else:
-                if not descriptor.key:
-                    if len(descriptor.keys) > 1:
-                        raise ValueError("not single-sig")
-                if self.key.xpub() != descriptor_xpubs[0]:
-                    raise ValueError("xpub does not match")
+            try:
+                self._validate_descriptor(descriptor, descriptor_xpubs)
+            except ValueError as e:
+                raise ValueError("Invalid Descriptor: %s" % e)
 
         self.wallet_data = wallet_data
         self.wallet_qr_format = qr_format

diff --git a/tests/pages/home_pages/test_pub_key_view.py b/tests/pages/home_pages/test_pub_key_view.py
@@ -78,7 +78,6 @@ def test_public_key(mocker, m5stickv, tdata):
 
         pub_key_viewer.public_key()
 
-        print(ctx.wallet.key.policy_type, ctx.wallet.key.script_type)
         version = "Zpub" if ctx.wallet.key.policy_type == TYPE_MULTISIG else "zpub"
         qr_view_calls = []
         print_qr_calls = []