add support for multiple keysets

Author: cblokh

fb_recover_keys.py

@@ -118,7 +118,7 @@ def main():
     for algo, info in privkeys.items():
         # info may be either None or tuple
         if info:
-            privkey, chaincode = info
+            privkey, chaincode, keyset_id = info
             pub = recover.get_public_key(algo, privkey)
             if show_xprv:
                 print(privkey_descriptions[algo] + ":\t" + recover.encode_extended_key(algo, privkey, chaincode, False))

test/test_recovery.py

@@ -8,7 +8,7 @@
 def test_recovery_basic():
     result = recover.restore_key_and_chaincode(TEST_DIR / "backup.zip", TEST_DIR / "priv.pem", "Thefireblocks1!")
 
-    ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
+    ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_ECDSA_SECP256K1']
     assert(ecdsa_priv_key == 0x473d1820ca4bf7cf6b018a8520b1ec0849cb99bce4fff45c5598723f67b3bd52)
     pub = recover.get_public_key("MPC_ECDSA_SECP256K1", ecdsa_priv_key)
     assert(pub == "021d84f3b6d7c6888f81c7cc381b658d85319f27e1ea9c93dff128667fb4b82ba0")
@@ -18,8 +18,8 @@ def test_recovery_basic():
 
 def test_full_recovery():
     result = recover.restore_key_and_chaincode(TEST_DIR / "backup_new.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!")
-    ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
-    eddsa_priv_key, eddsa_chaincode = result['MPC_EDDSA_ED25519']
+    ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_ECDSA_SECP256K1']
+    eddsa_priv_key, eddsa_chaincode, _ = result['MPC_EDDSA_ED25519']
 
     assert(ecdsa_priv_key == 0x66b1baf063db6e7152480334ebab0ab098e85f682b784754e46c18c962a1aa9d)
     assert(eddsa_priv_key == 0xd74820d02cc2aa09e2d0bcb36aeb92625b3d92c8d202063eab5513fd4453a44)
@@ -39,7 +39,7 @@ def test_full_recovery():
 
 def test_recovery_old_format():
     result = recover.restore_key_and_chaincode(TEST_DIR / "backup_old_format.zip", TEST_DIR / "priv.pem", "Thefireblocks1!")
-    ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
+    ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_ECDSA_SECP256K1']
 
     assert(ecdsa_priv_key == 0x473d1820ca4bf7cf6b018a8520b1ec0849cb99bce4fff45c5598723f67b3bd52)
     pub = recover.get_public_key("MPC_ECDSA_SECP256K1", ecdsa_priv_key)
@@ -50,8 +50,8 @@ def test_recovery_old_format():
 
 def test_cmp_recovery():
     result = recover.restore_key_and_chaincode(TEST_DIR / "backup_cmp.zip", TEST_DIR / "priv.pem", "Fireblocks1!")
-    ecdsa_priv_key, ecdsa_chaincode = result['MPC_CMP_ECDSA_SECP256K1']
-    eddsa_priv_key, eddsa_chaincode = result['MPC_CMP_EDDSA_ED25519']
+    ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_CMP_ECDSA_SECP256K1']
+    eddsa_priv_key, eddsa_chaincode, _ = result['MPC_CMP_EDDSA_ED25519']
 
     assert(ecdsa_priv_key == 0xf57c18e98a24ca0b36fbbd103233aff128b740426da189ce208545d44bbad050)
     assert(eddsa_priv_key == 0xa536dc2f2d744ae78eb26fdfb4b9e234a649525e0a1142bf900cd9c26987007)
@@ -77,8 +77,8 @@ def test_one_custom_chaincode_recovery():
     which encodes the chaincode. 
     '''
     result = recover.restore_key_and_chaincode(TEST_DIR / "backup_with_one_custom_chaincode.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!")
-    ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
-    eddsa_priv_key, eddsa_chaincode = result['MPC_EDDSA_ED25519']
+    ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_ECDSA_SECP256K1']
+    eddsa_priv_key, eddsa_chaincode, _ = result['MPC_EDDSA_ED25519']
 
     assert(ecdsa_chaincode != eddsa_chaincode)
     assert(ecdsa_priv_key == 0x66b1baf063db6e7152480334ebab0ab098e85f682b784754e46c18c962a1aa9d)
@@ -111,8 +111,8 @@ def test_two_custom_chaincode_recovery():
  
     '''
     result = recover.restore_key_and_chaincode(TEST_DIR / "backup_with_two_custom_chaincode.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!")
-    ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
-    eddsa_priv_key, eddsa_chaincode = result['MPC_EDDSA_ED25519']
+    ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_ECDSA_SECP256K1']
+    eddsa_priv_key, eddsa_chaincode, _ = result['MPC_EDDSA_ED25519']
 
     assert(ecdsa_chaincode != eddsa_chaincode)
     assert(ecdsa_priv_key == 0x66b1baf063db6e7152480334ebab0ab098e85f682b784754e46c18c962a1aa9d)
@@ -133,8 +133,8 @@ def test_two_custom_chaincode_recovery():
 
 def test_recovery_with_mobile_share_as_json():
     result = recover.restore_key_and_chaincode(TEST_DIR / "json_share.zip", TEST_DIR / "priv.pem", "Fireblocks1!")
-    ecdsa_priv_key, ecdsa_chaincode = result['MPC_CMP_ECDSA_SECP256K1']
-    eddsa_priv_key, eddsa_chaincode = result['MPC_EDDSA_ED25519']
+    ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_CMP_ECDSA_SECP256K1']
+    eddsa_priv_key, eddsa_chaincode, _ = result['MPC_EDDSA_ED25519']
 
     assert(ecdsa_priv_key == 0x83af98f2f2bdea33eb34177b311d89569725a401c1fc4d6046d266b1ca0dc382)
     assert(eddsa_priv_key == 0xd5c4d44f0f07aaa0bf18e039f28ec2131935ed696636f48ec46bab58e66296b)
@@ -154,8 +154,8 @@ def test_recovery_with_mobile_share_as_json():
 
 def test_recovery_with_master_keys():
     result = recover.restore_key_and_chaincode(TEST_DIR / "backup_with_master_key.zip", TEST_DIR / "priv_ncw.pem", "2#0Iw0Qov@&QP09p", key_pass="SECRET")
-    ecdsa_priv_key, ecdsa_chaincode = result["MPC_CMP_ECDSA_SECP256K1"]
-    eddsa_priv_key, eddsa_chaincode = result["MPC_EDDSA_ED25519"]
+    ecdsa_priv_key, ecdsa_chaincode, _ = result["MPC_CMP_ECDSA_SECP256K1"]
+    eddsa_priv_key, eddsa_chaincode, _ = result["MPC_EDDSA_ED25519"]
 
     pub = recover.get_public_key("MPC_ECDSA_SECP256K1", ecdsa_priv_key)
     assert pub == "033f5e4fb621e4cc777e5b9cdc0ef06c7b55042e9ce6c3bf013daab9fba29b37b8"

utils/metadata.py

@@ -17,6 +17,7 @@ class SigningKeyMetadata:
     public_key: str
     algorithm: str
     chain_code: bytes
+    keyset_id: int
 
 
 @dataclass
@@ -55,13 +56,18 @@ def parse_metadata_file(fp: IO[bytes]) -> RecoveryPackageMetadata:
         else:
             chain_code_for_this_key = default_chain_code
 
+        keyset_id_for_this_key = 1
+        if "keysetId" in key_metadata:
+            keyset_id_for_this_key = int(key_metadata["keysetId"])
+
         if len(chain_code_for_this_key) != 32:
             raise RecoveryErrorUnknownChainCode()
 
         signing_keys[key_id] = SigningKeyMetadata(
             public_key=metadata_public_key,
             algorithm=algo,
-            chain_code=chain_code_for_this_key
+            chain_code=chain_code_for_this_key,
+            keyset_id=keyset_id_for_this_key
         )
 
     master_keys_in_backup = obj.get("masterKeys", {})

utils/recover.py

@@ -239,7 +239,7 @@ def restore_key_and_chaincode(zip_path, private_pem_path, passphrase, key_pass=N
             print(f"Failed to recover {algo} key, expected public key is: {pub_from_metadata} calculated public key is: {pubkey_str}")
             privkeys[algo] = None
         else:
-            privkeys[algo] = privkey, chain_code_for_this_key
+            privkeys[algo] = privkey, chain_code_for_this_key, signing_keys[key_id].keyset_id
 
     if len(privkeys) == 0:
         raise RecoveryErrorPublicKeyNoMatch()