bip32: unserialize and store the network prefix

Author: darosior

bip32/bip32.py

@@ -183,10 +183,11 @@ def from_xpriv(cls, xpriv):
         :param xpriv: (str) The encoded serialized extended private key.
         """
         extended_key = base58.b58decode_check(xpriv)
-        (prefix, depth, fingerprint,
+        (network, depth, fingerprint,
          index, chaincode, key) = _unserialize_extended_key(extended_key)
         # We need to remove the trailing `0` before the actual private key !!
-        return BIP32(chaincode, key[1:], None, fingerprint, depth, index)
+        return BIP32(chaincode, key[1:], None, fingerprint, depth, index,
+                     network)
 
     @classmethod
     def from_xpub(cls, xpub):
@@ -195,9 +196,9 @@ def from_xpub(cls, xpub):
         :param xpub: (str) The encoded serialized extended public key.
         """
         extended_key = base58.b58decode_check(xpub)
-        (prefix, depth, fingerprint,
+        (network, depth, fingerprint,
          index, chaincode, key) = _unserialize_extended_key(extended_key)
-        return BIP32(chaincode, None, key, fingerprint, depth, index)
+        return BIP32(chaincode, None, key, fingerprint, depth, index, network)
 
     @classmethod
     def from_seed(cls, seed, network="main"):

bip32/utils.py

@@ -157,16 +157,20 @@ def _unserialize_extended_key(extended_key):
 
     :param extended_key: The extended key to unserialize __as bytes__
 
-    :return: prefix (int), depth (int), fingerprint (bytes), index (int),
+    :return: network (str), depth (int), fingerprint (bytes), index (int),
              chaincode (bytes), key (bytes)
     """
     assert isinstance(extended_key, bytes) and len(extended_key) == 78
     prefix = int.from_bytes(extended_key[:4], "big")
+    if prefix in list(ENCODING_PREFIX["main"].values()):
+        network = "main"
+    else:
+        network = "test"
     depth = extended_key[4]
     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
+    return network, depth, fingerprint, index, chaincode, key
 
 
 def _hardened_index_in_path(path):

tests/test_bip32.py

@@ -124,3 +124,23 @@ def test_sanity_tests():
     assert bip32.get_master_xpriv() == xpriv2
     assert bip32.get_xpriv_from_path([HARDENED_INDEX, 18]) == xpriv
     assert bip32.get_xpub_from_path([HARDENED_INDEX, 18]) == xpub
+    # We should recognize the networks..
+    # .. for xprivs:
+    bip32 = BIP32.from_xpriv("xprv9wHokC2KXdTSpEepFcu53hMDUHYfAtTaLEJEMyxBPAMf78hJg17WhL5FyeDUQH5KWmGjGgEb2j74gsZqgupWpPbZgP6uFmP8MYEy5BNbyET")
+    assert bip32.network == "main"
+    bip32 = BIP32.from_xpriv("tprv8ZgxMBicQKsPeCBsMzQCCb5JcW4S49MVL3EwhdZMF1RF71rgisZU4ZRvrHX6PZQEiNUABDLvYqpx8Lsccq8aGGR59qHAoLoE3iXYuDa8JTP")
+    assert bip32.network == "test"
+    # .. for xpubs:
+    bip32 = BIP32.from_xpub("xpub6AHA9hZDN11k2ijHMeS5QqHx2KP9aMBRhTDqANMnwVtdyw2TDYRmF8PjpvwUFcL1Et8Hj59S3gTSMcUQ5gAqTz3Wd8EsMTmF3DChhqPQBnU")
+    assert bip32.network == "main"
+    bip32 = BIP32.from_xpub("tpubD6NzVbkrYhZ4WN3WiKRjeo2eGyYNiKNg8vcQ1UjLNJJaDvoFhmR1XwJsbo5S4vicSPoWQBThR3Rt8grXtP47c1AnoiXMrEmFdRZupxJzH1j")
+    assert bip32.network == "test"
+    # We should create valid network encoding..
+    assert BIP32.from_seed(os.urandom(32),
+                           "test").get_master_xpub().startswith("tpub")
+    assert BIP32.from_seed(os.urandom(32),
+                           "test").get_master_xpriv().startswith("tprv")
+    assert BIP32.from_seed(os.urandom(32),
+                           "main").get_master_xpub().startswith("xpub")
+    assert BIP32.from_seed(os.urandom(32),
+                           "main").get_master_xpriv().startswith("xprv")