From d2f75b7da5b4d2c300ae52ecfae7cefe86204b6c Mon Sep 17 00:00:00 2001
From: SomberNight <somber.night@protonmail.com>
Date: Sun, 26 Feb 2023 11:40:42 +0000
Subject: [PATCH] descriptor.py: don't allow ypub/zpub inside descriptors

---
 electrum/bip32.py              | 10 +++++++++-
 electrum/descriptor.py         |  4 ++--
 electrum/tests/test_bitcoin.py | 22 ++++++++++++++++++++++
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/electrum/bip32.py b/electrum/bip32.py
index f27952e0d..74ee1aeb8 100644
--- a/electrum/bip32.py
+++ b/electrum/bip32.py
@@ -124,7 +124,13 @@ class BIP32Node(NamedTuple):
     child_number: bytes = b'\x00'*4
 
     @classmethod
-    def from_xkey(cls, xkey: str, *, net=None) -> 'BIP32Node':
+    def from_xkey(
+        cls,
+        xkey: str,
+        *,
+        net=None,
+        allow_custom_headers: bool = True,  # to also accept ypub/zpub
+    ) -> 'BIP32Node':
         if net is None:
             net = constants.net
         xkey = DecodeBase58Check(xkey)
@@ -145,6 +151,8 @@ class BIP32Node(NamedTuple):
         else:
             raise InvalidMasterKeyVersionBytes(f'Invalid extended key format: {hex(header)}')
         xtype = headers_inv[header]
+        if not allow_custom_headers and xtype != "standard":
+            raise ValueError(f"only standard xpub/xprv allowed. found custom xtype={xtype}")
         if is_private:
             eckey = ecc.ECPrivkey(xkey[13 + 33:])
         else:
diff --git a/electrum/descriptor.py b/electrum/descriptor.py
index fa7aa1354..bacfa71cc 100644
--- a/electrum/descriptor.py
+++ b/electrum/descriptor.py
@@ -121,8 +121,8 @@ class PubkeyProvider(object):
             unhexlify(self.pubkey)
             # Is hex, normal pubkey
         except Exception:
-            # Not hex, maybe xpub
-            self.extkey = BIP32Node.from_xkey(pubkey)
+            # Not hex, maybe xpub (but don't allow ypub/zpub)
+            self.extkey = BIP32Node.from_xkey(pubkey, allow_custom_headers=False)
 
     @classmethod
     def parse(cls, s: str) -> 'PubkeyProvider':
diff --git a/electrum/tests/test_bitcoin.py b/electrum/tests/test_bitcoin.py
index 08e881ecd..c3026a0d2 100644
--- a/electrum/tests/test_bitcoin.py
+++ b/electrum/tests/test_bitcoin.py
@@ -815,6 +815,28 @@ class Test_xprv_xpub(ElectrumTestCase):
         self.assertFalse(is_xprv('xprv1nval1d'))
         self.assertFalse(is_xprv('xprv661MyMwAqRbcFWohJWt7PHsFEJfZAvw9ZxwQoDa4SoMgsDDM1T7WK3u9E4edkC4ugRnZ8E4xDZRpk8Rnts3Nbt97dPwT52WRONGBADWRONG'))
 
+    def test_bip32_from_xkey(self):
+        bip32node1 = BIP32Node.from_xkey("xpub6H1LXWLaKsWFhvm6RVpEL9P4KfRZSW7abD2ttkWP3SSQvnyA8FSVqNTEcYFgJS2UaFcxupHiYkro49S8yGasTvXEYBVPamhGW6cFJodrTHy")
+        self.assertEqual(
+            BIP32Node(
+                xtype='standard',
+                eckey=ecc.ECPubkey(bytes.fromhex("022a471424da5e657499d1ff51cb43c47481a03b1e77f951fe64cec9f5a48f7011")),
+                chaincode=bytes.fromhex("c783e67b921d2beb8f6b389cc646d7263b4145701dadd2161548a8b078e65e9e"),
+                depth=5,
+                fingerprint=bytes.fromhex("d880d7d8"),
+                child_number=bytes.fromhex("3b9aca00"),
+            ),
+            bip32node1)
+        with self.assertRaises(ValueError):
+            BIP32Node.from_xkey(
+                "zpub6jftahH18ngZyLeqfLBFAm7YaWFVttE9pku5pNMX2qPzTjoq1FVgZMmhjecyB2nqFb31gHE9vNvbaggU6vvWpNZbXEWLLUjYjFqG95LNyT8",
+                allow_custom_headers=False)
+        bip32node2 = BIP32Node.from_xkey(
+            "zpub6jftahH18ngZyLeqfLBFAm7YaWFVttE9pku5pNMX2qPzTjoq1FVgZMmhjecyB2nqFb31gHE9vNvbaggU6vvWpNZbXEWLLUjYjFqG95LNyT8",
+            allow_custom_headers=True)
+        self.assertEqual(bytes.fromhex("03f18e53f3386a5f9a9d2c369ad3b84b429eb397b4bc69ce600f2d833b54ba32f4"),
+                         bip32node2.eckey.get_public_key_bytes(compressed=True))
+
     def test_is_bip32_derivation(self):
         self.assertTrue(is_bip32_derivation("m/0'/1"))
         self.assertTrue(is_bip32_derivation("m/0'/0'"))