bip32: fixup serialization

We wouldn't save the position of our master node in a potential already-existing tree.

Author: darosior

bip32/bip32.py

@@ -11,7 +11,8 @@
 
 
 class BIP32:
-    def __init__(self, chaincode, privkey=None, pubkey=None):
+    def __init__(self, chaincode, privkey=None, pubkey=None, fingerprint=None,
+                 depth=0, index=0):
         """
         :param chaincode: The master chaincode, used to derive keys. As bytes.
         :param privkey: The master private key for this index (default 0).
@@ -20,6 +21,13 @@ def __init__(self, chaincode, privkey=None, pubkey=None):
         :param pubkey: The master public key for this index (default 0).
                        Can be None if private key is specified.
                        Compressed format. As bytes.
+        :param fingeprint: If we are instanciated from an xpub/xpriv, we need
+                           to remember the parent's pubkey fingerprint to
+                           reserialize !
+        :param depth: If we are instanciated from an existing extended key, we
+                      need this for serialization.
+        :param index: If we are instanciated from an existing extended key, we
+                      need this for serialization.
         """
         assert isinstance(chaincode, bytes)
         assert privkey is not None or pubkey is not None
@@ -32,6 +40,9 @@ def __init__(self, chaincode, privkey=None, pubkey=None):
         self.master_chaincode = chaincode
         self.master_privkey = privkey
         self.master_pubkey = pubkey
+        self.parent_fingerprint = fingerprint
+        self.depth = depth
+        self.index = index
 
     def get_extended_privkey_from_path(self, path):
         """Get an extended privkey from a list of indexes (path).
@@ -111,7 +122,7 @@ def get_xpriv_from_path(self, path):
         else:
             parent_pubkey = self.get_pubkey_from_path(path[:-1])
         chaincode, privkey = self.get_extended_privkey_from_path(path)
-        extended_key = _serialize_extended_key(privkey, len(path),
+        extended_key = _serialize_extended_key(privkey, self.depth + len(path),
                                                parent_pubkey,
                                                path[-1], chaincode)
         return base58.b58encode_check(extended_key).decode()
@@ -130,21 +141,27 @@ def get_xpub_from_path(self, path):
         else:
             parent_pubkey = self.get_pubkey_from_path(path[:-1])
         chaincode, pubkey = self.get_extended_pubkey_from_path(path)
-        extended_key = _serialize_extended_key(pubkey, len(path),
+        extended_key = _serialize_extended_key(pubkey, self.depth + len(path),
                                                parent_pubkey,
                                                path[-1], chaincode)
         return base58.b58encode_check(extended_key).decode()
 
     def get_master_xpriv(self):
         """Get the encoded extended private key of the master private key"""
-        extended_key = _serialize_extended_key(self.master_privkey, 0,
-                                               None, 0, self.master_chaincode)
+        extended_key = _serialize_extended_key(self.master_privkey, self.depth,
+                                               self.parent_fingerprint,
+                                               self.index,
+                                               self.master_chaincode)
+        print(extended_key)
+        print(base58.b58encode_check(extended_key))
         return base58.b58encode_check(extended_key).decode()
 
     def get_master_xpub(self):
         """Get the encoded extended public key of the master public key"""
-        extended_key = _serialize_extended_key(self.master_pubkey, 0,
-                                               None, 0, self.master_chaincode)
+        extended_key = _serialize_extended_key(self.master_pubkey, self.depth,
+                                               self.parent_fingerprint,
+                                               self.index,
+                                               self.master_chaincode)
         return base58.b58encode_check(extended_key).decode()
 
     @classmethod
@@ -153,11 +170,16 @@ def from_xpriv(cls, xpriv):
 
         :param xpriv: (str) The encoded serialized extended private key.
         """
+        print(xpriv)
         extended_key = base58.b58decode_check(xpriv)
         (prefix, depth, fingerprint,
          index, chaincode, key) = _unserialize_extended_key(extended_key)
+        serialized = _serialize_extended_key(key[1:], depth, fingerprint, index,
+                                             chaincode)
+        print(extended_key, serialized, extended_key == serialized)
+        print(base58.b58encode_check(extended_key), base58.b58encode_check(serialized), extended_key == serialized)
         # We need to remove the trailing `0` before the actual private key !!
-        return BIP32(chaincode, privkey=key[1:], pubkey=None)
+        return BIP32(chaincode, key[1:], None, fingerprint, depth, index)
 
     @classmethod
     def from_xpub(cls, xpub):
@@ -168,7 +190,7 @@ def from_xpub(cls, xpub):
         extended_key = base58.b58decode_check(xpub)
         (prefix, depth, fingerprint,
          index, chaincode, key) = _unserialize_extended_key(extended_key)
-        return BIP32(chaincode, privkey=None, pubkey=key)
+        return BIP32(chaincode, None, key, fingerprint, depth, index)
 
     @classmethod
     def from_seed(cls, seed):
@@ -178,4 +200,4 @@ def from_seed(cls, seed):
         """
         secret = hmac.new("Bitcoin seed".encode(), seed,
                           hashlib.sha512).digest()
-        return BIP32(secret[32:], secret[:32], None)
+        return BIP32(secret[32:], secret[:32])

bip32/utils.py

@@ -104,15 +104,15 @@ def _pubkey_to_fingerprint(pubkey):
     return rip.digest()[:4]
 
 
-def _serialize_extended_key(key, depth, parent_pubkey, index, chaincode,
+def _serialize_extended_key(key, depth, parent, index, chaincode,
                             network="main"):
     """Serialize an extended private *OR* public key, as spec by bip-0032.
 
     :param key: The public or private key to serialize. Note that if this is
                 a public key it MUST be compressed.
     :param depth: 0x00 for master nodes, 0x01 for level-1 derived keys, etc..
-    :param parent_pubkey: The parent pubkey used to derive the fingerprint.
-                          None if master.
+    :param parent: The parent pubkey used to derive the fingerprint, or the
+                   fingerprint itself None if master.
     :param index: The index of the key being serialized. 0x00000000 if master.
     :param chaincode: The chain code (not the labs !!).
 
@@ -122,9 +122,15 @@ def _serialize_extended_key(key, depth, parent_pubkey, index, chaincode,
         assert isinstance(param, bytes)
     for param in {depth, index}:
         assert isinstance(param, int)
-    if parent_pubkey:
-        assert isinstance(parent_pubkey, bytes) and len(parent_pubkey) == 33
-        fingerprint = _pubkey_to_fingerprint(parent_pubkey)
+    if parent:
+        assert isinstance(parent, bytes)
+        if len(parent) == 33:
+            fingerprint = _pubkey_to_fingerprint(parent)
+        elif len(parent) == 4:
+            fingerprint = parent
+        else:
+            raise ValueError("Bad parent, a fingerprint or a pubkey is"
+                             " required")
     else:
         fingerprint = bytes(4)  # master
     # A privkey or a compressed pubkey
@@ -155,12 +161,11 @@ def _unserialize_extended_key(extended_key):
     assert isinstance(extended_key, bytes) and len(extended_key) == 78
     prefix = int.from_bytes(extended_key[:4], "big")
     depth = extended_key[4]
-    fingerprint = extended_key[4:9]
+    fingerprint = extended_key[5:9]
     index = int.from_bytes(extended_key[9:13], "big")
     chaincode, key = extended_key[13:45], extended_key[45:]
     return prefix, depth, fingerprint, index, chaincode, key
 
 
 def _hardened_index_in_path(path):
     return len([i for i in path if i & HARDENED_INDEX]) > 0
-

tests/test_bip32.py

@@ -94,5 +94,8 @@ def test_sanity_tests():
     assert xpub == "xpub6FXqzrT8Uh8gs81aM7xMJPzAEacxEUg7aC8yA4xz4xmPfEoKxJdVT9Hdwm3LwVQrSos2rhGDt8aGGHvdLr5LLAjK8pXFkbSpzGoGTXjd4z9"
     # Now if we our master is m/0'/0'/14, we should derive the same keys for
     # m/0'/18 !
-    bip32 = BIP32.from_xpriv(bip32.get_xpriv_from_path([HARDENED_INDEX,
-                                                        HARDENED_INDEX, 14]))
+    xpriv2 = bip32.get_xpriv_from_path([HARDENED_INDEX, HARDENED_INDEX, 14])
+    assert xpriv2 == "xprv9yQJmvQMywM5i7UNuZ4RQ1A9rEMwAJCExPardkmBCB46S3vBqNEatSwLUrwLNLHBu1Kd9aGxGKDD5YAfs6hRzpYthciAHjtGadxgV2PeqY9"
+    bip32 = BIP32.from_xpriv(xpriv2)
+    assert bip32.get_master_xpriv() == xpriv2
+    assert bip32.get_xpriv_from_path([HARDENED_INDEX, 18]) == xpriv