From 9ab231520ef7de96ed333555ff69d448ba4ea01f Mon Sep 17 00:00:00 2001
From: Kristaps Kaupe <kristaps@blogiem.lv>
Date: Thu, 17 Feb 2022 13:16:01 +0200
Subject: [PATCH] Stricter BIP32 decoding and test vector 5

---
 .../jmbitcoin/secp256k1_deterministic.py      | 15 ++++++-
 jmbitcoin/test/test_bip32.py                  | 44 +++++++++++++++----
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/jmbitcoin/jmbitcoin/secp256k1_deterministic.py b/jmbitcoin/jmbitcoin/secp256k1_deterministic.py
index 6835afd..eb4fa3d 100644
--- a/jmbitcoin/jmbitcoin/secp256k1_deterministic.py
+++ b/jmbitcoin/jmbitcoin/secp256k1_deterministic.py
@@ -61,12 +61,23 @@ def bip32_deserialize(data):
     if Hash(dbin[:-4])[:4] != dbin[-4:]:
         raise Exception("Invalid checksum")
     vbytes = dbin[0:4]
+    if vbytes not in PRIVATE and vbytes not in PUBLIC:
+        raise Exception("Invalid vbytes {}".format(vbytes))
     depth = dbin[4]
     fingerprint = dbin[5:9]
-    i = struct.unpack(b'>L',dbin[9:13])[0]
+    child_num = struct.unpack(b'>L',dbin[9:13])[0]
+    if depth == 0 and (fingerprint != b'\x00'*4 or child_num != 0):
+        raise Exception("Invalid master key, depth = {}, fingerprint = {}, child_num = {}".format(depth, fingerprint, child_num))
     chaincode = dbin[13:45]
     key = dbin[46:78] + b'\x01' if vbytes in PRIVATE else dbin[45:78]
-    return (vbytes, depth, fingerprint, i, chaincode, key)
+    if vbytes in PUBLIC and not is_valid_pubkey(key):
+        raise Exception("Invalid public key")
+    if vbytes in PRIVATE:
+        if dbin[45] != 0:
+            raise Exception("Invalid private key")
+        # check for validity, this will raise exception
+        privkey_to_pubkey(key)
+    return (vbytes, depth, fingerprint, child_num, chaincode, key)
 
 def raw_bip32_privtopub(rawtuple):
     vbytes, depth, fingerprint, i, chaincode, key = rawtuple
diff --git a/jmbitcoin/test/test_bip32.py b/jmbitcoin/test/test_bip32.py
index e26420e..2d7ed95 100644
--- a/jmbitcoin/test/test_bip32.py
+++ b/jmbitcoin/test/test_bip32.py
@@ -95,15 +95,41 @@ def test_bip32_vector(vector):
         assert pub == vector['keys'][i][
             1], 'failed: child pub key, should be: ' + vector['keys'][i][1]
 
-def test_invalid_bip32_key():
-    #keys with invalid checksum
-    bad_keys = (
-    'xprv9uHRZZhk6KAJC1avXpDAp4MDc3sQKNxDiPvvkX8BrrngLNv1TxvUxt4cV1rGL5hj6KCesnDYUhd7oWgT11eZG7XnxHrnYeSvkzY7d2bhkJ7',
-    'xpub68Gmy5EdvgibQVfPdqkBBCHxA5htiqg55crXYuXoQQKfDBFA1WEjWgP6LHhwBZeNK1VTsfTFUHCdrfp1bgwQ9xv5ski8PX9rL2dZXvgGDnw'
-    )
-    for x in bad_keys:
-        with pytest.raises(Exception) as e_info:
-            fake_tuple = btc.bip32_deserialize(x)
+@pytest.mark.parametrize(
+    'bad_key',
+    [
+        #keys with invalid checksum
+        'xprv9uHRZZhk6KAJC1avXpDAp4MDc3sQKNxDiPvvkX8BrrngLNv1TxvUxt4cV1rGL5hj6KCesnDYUhd7oWgT11eZG7XnxHrnYeSvkzY7d2bhkJ7',
+        'xpub68Gmy5EdvgibQVfPdqkBBCHxA5htiqg55crXYuXoQQKfDBFA1WEjWgP6LHhwBZeNK1VTsfTFUHCdrfp1bgwQ9xv5ski8PX9rL2dZXvgGDnw',
+        ## BIP32 test vector 5
+        'xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHL',
+        #pubkey version / prvkey mismatch
+        'xpub661MyMwAqRbcEYS8w7XLSVeEsBXy79zSzH1J8vCdxAZningWLdN3zgtU6LBpB85b3D2yc8sfvZU521AAwdZafEz7mnzBBsz4wKY5fTtTQBm',
+        #prvkey version / pubkey mismatch
+        'xprv9s21ZrQH143K24Mfq5zL5MhWK9hUhhGbd45hLXo2Pq2oqzMMo63oStZzFGTQQD3dC4H2D5GBj7vWvSQaaBv5cxi9gafk7NF3pnBju6dwKvH',
+        #keys with invalid prefix
+        'xpub661MyMwAqRbcEYS8w7XLSVeEsBXy79zSzH1J8vCdxAZningWLdN3zgtU6Txnt3siSujt9RCVYsx4qHZGc62TG4McvMGcAUjeuwZdduYEvFn',
+        'xprv9s21ZrQH143K24Mfq5zL5MhWK9hUhhGbd45hLXo2Pq2oqzMMo63oStZzFGpWnsj83BHtEy5Zt8CcDr1UiRXuWCmTQLxEK9vbz5gPstX92JQ',
+        'xpub661MyMwAqRbcEYS8w7XLSVeEsBXy79zSzH1J8vCdxAZningWLdN3zgtU6N8ZMMXctdiCjxTNq964yKkwrkBJJwpzZS4HS2fxvyYUA4q2Xe4',
+        'xprv9s21ZrQH143K24Mfq5zL5MhWK9hUhhGbd45hLXo2Pq2oqzMMo63oStZzFAzHGBP2UuGCqWLTAPLcMtD9y5gkZ6Eq3Rjuahrv17fEQ3Qen6J',
+        #zero depth with non-zero parent fingerprint
+        'xprv9s2SPatNQ9Vc6GTbVMFPFo7jsaZySyzk7L8n2uqKXJen3KUmvQNTuLh3fhZMBoG3G4ZW1N2kZuHEPY53qmbZzCHshoQnNf4GvELZfqTUrcv',
+        'xpub661no6RGEX3uJkY4bNnPcw4URcQTrSibUZ4NqJEw5eBkv7ovTwgiT91XX27VbEXGENhYRCf7hyEbWrR3FewATdCEebj6znwMfQkhRYHRLpJ',
+        #zero depth with non-zero index
+        'xprv9s21ZrQH4r4TsiLvyLXqM9P7k1K3EYhA1kkD6xuquB5i39AU8KF42acDyL3qsDbU9NmZn6MsGSUYZEsuoePmjzsB3eFKSUEh3Gu1N3cqVUN',
+        'xpub661MyMwAuDcm6CRQ5N4qiHKrJ39Xe1R1NyfouMKTTWcguwVcfrZJaNvhpebzGerh7gucBvzEQWRugZDuDXjNDRmXzSZe4c7mnTK97pTvGS8'
+        #unknown extended key version
+        'DMwo58pR1QLEFihHiXPVykYB6fJmsTeHvyTp7hRThAtCX8CvYzgPcn8XnmdfHGMQzT7ayAmfo4z3gY5KfbrZWZ6St24UVf2Qgo6oujFktLHdHY4',
+        'DMwo58pR1QLEFihHiXPVykYB6fJmsTeHvyTp7hRThAtCX8CvYzgPcn8XnmdfHPmHJiEDXkTiJTVV9rHEBUem2mwVbbNfvT2MTcAqj3nesx8uBf9',
+        #private key 0 not in 1..n-1
+        'xprv9s21ZrQH143K24Mfq5zL5MhWK9hUhhGbd45hLXo2Pq2oqzMMo63oStZzF93Y5wvzdUayhgkkFoicQZcP3y52uPPxFnfoLZB21Teqt1VvEHx',
+        'xprv9s21ZrQH143K24Mfq5zL5MhWK9hUhhGbd45hLXo2Pq2oqzMMo63oStZzFAzHGBP2UuGCqWLTAPLcMtD5SDKr24z3aiUvKr9bJpdrcLg1y3G',
+        #invalid pubkey 020000000000000000000000000000000000000000000000000000000000000007
+        'xpub661MyMwAqRbcEYS8w7XLSVeEsBXy79zSzH1J8vCdxAZningWLdN3zgtU6Q5JXayek4PRsn35jii4veMimro1xefsM58PgBMrvdYre8QyULY',
+    ])
+def test_invalid_bip32_key(bad_key):
+    with pytest.raises(Exception) as e_info:
+        fake_tuple = btc.bip32_deserialize(bad_key)
 
 def test_ckd_pubkeys():
     pub = 'xpub68Gmy5EdvgibQVfPdqkBBCHxA5htiqg55crXYuXoQRKfDBFA1WEjWgP6LHhwBZeNK1VTsfTFUHCdrfp1bgwQ9xv5ski8PX9rL2dZXvgGDnw'