Skip to content

Commit 147b81f

Browse files
cblokhcblokh
committed
add support for multiple keysets
1 parent 767ed51 commit 147b81f

File tree

4 files changed

+59
-30
lines changed

4 files changed

+59
-30
lines changed

‎fb_recover_keys.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -118,13 +118,14 @@ def main():
118118
for algo, info in privkeys.items():
119119
# info may be either None or tuple
120120
if info:
121-
privkey, chaincode = info
121+
privkey, chaincode, keyset_id = info
122+
keyset_id_str = " [" + keyset_id + "]"
122123
pub = recover.get_public_key(algo, privkey)
123124
if show_xprv:
124-
print(privkey_descriptions[algo] + ":\t" + recover.encode_extended_key(algo, privkey, chaincode, False))
125-
print(pubkey_descriptions[algo] + ":\t%s\t%s" % (recover.encode_extended_key(algo, pub, chaincode, True), colored("Verified!","green")))
125+
print(privkey_descriptions[algo] + keyset_id_str + ":\t" + recover.encode_extended_key(algo, privkey, chaincode, False))
126+
print(pubkey_descriptions[algo] + keyset_id_str + ":\t%s\t%s" % (recover.encode_extended_key(algo, pub, chaincode, True), colored("Verified!","green")))
126127
else:
127-
print(pubkey_descriptions[algo] + ":\t%s" % (colored("Verification failed","red")))
128+
print(pubkey_descriptions[algo] + keyset_id_str + ":\t%s" % (colored("Verification failed","red")))
128129

129130
if __name__ == "__main__" :
130131
main()

‎test/test_recovery.py‎

Lines changed: 26 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -6,20 +6,21 @@
66

77

88
def test_recovery_basic():
9-
result = recover.restore_key_and_chaincode(TEST_DIR / "backup.zip", TEST_DIR / "priv.pem", "Thefireblocks1!")
10-
11-
ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
9+
result = recover.strip_keyset_from_results(recover.restore_key_and_chaincode(TEST_DIR / "backup.zip", TEST_DIR / "priv.pem", "Thefireblocks1!"))
10+
print(result)
11+
ecdsa_priv_key, ecdsa_chaincode, keyset_id = result['MPC_ECDSA_SECP256K1']
1212
assert(ecdsa_priv_key == 0x473d1820ca4bf7cf6b018a8520b1ec0849cb99bce4fff45c5598723f67b3bd52)
1313
pub = recover.get_public_key("MPC_ECDSA_SECP256K1", ecdsa_priv_key)
1414
assert(pub == "021d84f3b6d7c6888f81c7cc381b658d85319f27e1ea9c93dff128667fb4b82ba0")
1515
assert(recover.encode_extended_key('MPC_ECDSA_SECP256K1', ecdsa_priv_key, ecdsa_chaincode, False) == "xprv9s21ZrQH143K24Mfq5zL5MhWK9hUhhGbd45hLXo2Pq2oqzMMo63oStZzF9aunJDs4SsrmoxycAo6xxBTHawSz5sYxEy8TpCkv66Sci373DJ")
1616
assert(recover.encode_extended_key('MPC_ECDSA_SECP256K1', pub, ecdsa_chaincode, True) == "xpub661MyMwAqRbcEYS8w7XLSVeEsBXy79zSzH1J8vCdxAZningWLdN3zgtU6QJJZSgiCXT6sq7wa2jCk5t4Vv1r1E4q1venKghAAdyzieufGyX")
17+
assert(keyset_id == 1)
1718

1819

1920
def test_full_recovery():
20-
result = recover.restore_key_and_chaincode(TEST_DIR / "backup_new.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!")
21-
ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
22-
eddsa_priv_key, eddsa_chaincode = result['MPC_EDDSA_ED25519']
21+
result = recover.strip_keyset_from_results(recover.restore_key_and_chaincode(TEST_DIR / "backup_new.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!"))
22+
ecdsa_priv_key, ecdsa_chaincode, keyset_id = result['MPC_ECDSA_SECP256K1']
23+
eddsa_priv_key, eddsa_chaincode, keyset_id_eddsa = result['MPC_EDDSA_ED25519']
2324

2425
assert(ecdsa_priv_key == 0x66b1baf063db6e7152480334ebab0ab098e85f682b784754e46c18c962a1aa9d)
2526
assert(eddsa_priv_key == 0xd74820d02cc2aa09e2d0bcb36aeb92625b3d92c8d202063eab5513fd4453a44)
@@ -35,11 +36,13 @@ def test_full_recovery():
3536
assert(pub == "0050cfee85dabebed78f43e94a1b7afd13c20461ad66efa083779bdeffd22269d9")
3637
assert(recover.encode_extended_key('MPC_EDDSA_ED25519', eddsa_priv_key, eddsa_chaincode, False) == "fprv4LsXPWzhTTp9ax8NGVwbnRFuT3avVQ4ydHNWcu8hCGZd18TRKxgAzbrpY9bLJRe4Y2AyX9TfQdDPbmqEYoDCTju9QFZbUgdsxsmUgfvuEDK")
3738
assert(recover.encode_extended_key('MPC_EDDSA_ED25519', pub, eddsa_chaincode, True) == "fpub8sZZXw2wbqVpURAAA9cCBpv2256rejFtCayHuRAzcYN1qciBxMVmB6UgiDAQTUZh5EP9JZciPQPjKAHyqPYHELqEHWkvo1sxreEJgLyfCJj")
39+
assert(keyset_id == 1)
40+
assert(keyset_id_eddsa == 1)
3841

3942

4043
def test_recovery_old_format():
41-
result = recover.restore_key_and_chaincode(TEST_DIR / "backup_old_format.zip", TEST_DIR / "priv.pem", "Thefireblocks1!")
42-
ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
44+
result = recover.strip_keyset_from_results(recover.restore_key_and_chaincode(TEST_DIR / "backup_old_format.zip", TEST_DIR / "priv.pem", "Thefireblocks1!"))
45+
ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_ECDSA_SECP256K1']
4346

4447
assert(ecdsa_priv_key == 0x473d1820ca4bf7cf6b018a8520b1ec0849cb99bce4fff45c5598723f67b3bd52)
4548
pub = recover.get_public_key("MPC_ECDSA_SECP256K1", ecdsa_priv_key)
@@ -49,9 +52,9 @@ def test_recovery_old_format():
4952

5053

5154
def test_cmp_recovery():
52-
result = recover.restore_key_and_chaincode(TEST_DIR / "backup_cmp.zip", TEST_DIR / "priv.pem", "Fireblocks1!")
53-
ecdsa_priv_key, ecdsa_chaincode = result['MPC_CMP_ECDSA_SECP256K1']
54-
eddsa_priv_key, eddsa_chaincode = result['MPC_CMP_EDDSA_ED25519']
55+
result = recover.strip_keyset_from_results(recover.restore_key_and_chaincode(TEST_DIR / "backup_cmp.zip", TEST_DIR / "priv.pem", "Fireblocks1!"))
56+
ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_CMP_ECDSA_SECP256K1']
57+
eddsa_priv_key, eddsa_chaincode, _ = result['MPC_CMP_EDDSA_ED25519']
5558

5659
assert(ecdsa_priv_key == 0xf57c18e98a24ca0b36fbbd103233aff128b740426da189ce208545d44bbad050)
5760
assert(eddsa_priv_key == 0xa536dc2f2d744ae78eb26fdfb4b9e234a649525e0a1142bf900cd9c26987007)
@@ -76,9 +79,9 @@ def test_one_custom_chaincode_recovery():
7679
Hence all the extracted keys are they same, and differce lies mostly in the extended form of the key,
7780
which encodes the chaincode.
7881
'''
79-
result = recover.restore_key_and_chaincode(TEST_DIR / "backup_with_one_custom_chaincode.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!")
80-
ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
81-
eddsa_priv_key, eddsa_chaincode = result['MPC_EDDSA_ED25519']
82+
result = recover.strip_keyset_from_results(recover.restore_key_and_chaincode(TEST_DIR / "backup_with_one_custom_chaincode.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!"))
83+
ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_ECDSA_SECP256K1']
84+
eddsa_priv_key, eddsa_chaincode, _ = result['MPC_EDDSA_ED25519']
8285

8386
assert(ecdsa_chaincode != eddsa_chaincode)
8487
assert(ecdsa_priv_key == 0x66b1baf063db6e7152480334ebab0ab098e85f682b784754e46c18c962a1aa9d)
@@ -110,9 +113,9 @@ def test_two_custom_chaincode_recovery():
110113
only the extended forms of the keys are different, as they encode the respective chaincodes.
111114
112115
'''
113-
result = recover.restore_key_and_chaincode(TEST_DIR / "backup_with_two_custom_chaincode.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!")
114-
ecdsa_priv_key, ecdsa_chaincode = result['MPC_ECDSA_SECP256K1']
115-
eddsa_priv_key, eddsa_chaincode = result['MPC_EDDSA_ED25519']
116+
result = recover.strip_keyset_from_results(recover.restore_key_and_chaincode(TEST_DIR / "backup_with_two_custom_chaincode.zip", TEST_DIR / "priv2.pem", "Thefireblocks1!"))
117+
ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_ECDSA_SECP256K1']
118+
eddsa_priv_key, eddsa_chaincode, _ = result['MPC_EDDSA_ED25519']
116119

117120
assert(ecdsa_chaincode != eddsa_chaincode)
118121
assert(ecdsa_priv_key == 0x66b1baf063db6e7152480334ebab0ab098e85f682b784754e46c18c962a1aa9d)
@@ -132,9 +135,9 @@ def test_two_custom_chaincode_recovery():
132135

133136

134137
def test_recovery_with_mobile_share_as_json():
135-
result = recover.restore_key_and_chaincode(TEST_DIR / "json_share.zip", TEST_DIR / "priv.pem", "Fireblocks1!")
136-
ecdsa_priv_key, ecdsa_chaincode = result['MPC_CMP_ECDSA_SECP256K1']
137-
eddsa_priv_key, eddsa_chaincode = result['MPC_EDDSA_ED25519']
138+
result = recover.strip_keyset_from_results(recover.restore_key_and_chaincode(TEST_DIR / "json_share.zip", TEST_DIR / "priv.pem", "Fireblocks1!"))
139+
ecdsa_priv_key, ecdsa_chaincode, _ = result['MPC_CMP_ECDSA_SECP256K1']
140+
eddsa_priv_key, eddsa_chaincode, _ = result['MPC_EDDSA_ED25519']
138141

139142
assert(ecdsa_priv_key == 0x83af98f2f2bdea33eb34177b311d89569725a401c1fc4d6046d266b1ca0dc382)
140143
assert(eddsa_priv_key == 0xd5c4d44f0f07aaa0bf18e039f28ec2131935ed696636f48ec46bab58e66296b)
@@ -153,9 +156,9 @@ def test_recovery_with_mobile_share_as_json():
153156

154157

155158
def test_recovery_with_master_keys():
156-
result = recover.restore_key_and_chaincode(TEST_DIR / "backup_with_master_key.zip", TEST_DIR / "priv_ncw.pem", "2#0Iw0Qov@&QP09p", key_pass="SECRET")
157-
ecdsa_priv_key, ecdsa_chaincode = result["MPC_CMP_ECDSA_SECP256K1"]
158-
eddsa_priv_key, eddsa_chaincode = result["MPC_EDDSA_ED25519"]
159+
result = recover.strip_keyset_from_results(recover.restore_key_and_chaincode(TEST_DIR / "backup_with_master_key.zip", TEST_DIR / "priv_ncw.pem", "2#0Iw0Qov@&QP09p", key_pass="SECRET"))
160+
ecdsa_priv_key, ecdsa_chaincode, _ = result["MPC_CMP_ECDSA_SECP256K1"]
161+
eddsa_priv_key, eddsa_chaincode, _ = result["MPC_EDDSA_ED25519"]
159162

160163
pub = recover.get_public_key("MPC_ECDSA_SECP256K1", ecdsa_priv_key)
161164
assert pub == "033f5e4fb621e4cc777e5b9cdc0ef06c7b55042e9ce6c3bf013daab9fba29b37b8"

‎utils/metadata.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ class SigningKeyMetadata:
1717
public_key: str
1818
algorithm: str
1919
chain_code: bytes
20+
keyset_id: int
2021

2122

2223
@dataclass
@@ -55,13 +56,18 @@ def parse_metadata_file(fp: IO[bytes]) -> RecoveryPackageMetadata:
5556
else:
5657
chain_code_for_this_key = default_chain_code
5758

59+
keyset_id_for_this_key = 1
60+
if "keysetId" in key_metadata:
61+
keyset_id_for_this_key = int(key_metadata["keysetId"])
62+
5863
if len(chain_code_for_this_key) != 32:
5964
raise RecoveryErrorUnknownChainCode()
6065

6166
signing_keys[key_id] = SigningKeyMetadata(
6267
public_key=metadata_public_key,
6368
algorithm=algo,
64-
chain_code=chain_code_for_this_key
69+
chain_code=chain_code_for_this_key,
70+
keyset_id=keyset_id_for_this_key
6571
)
6672

6773
master_keys_in_backup = obj.get("masterKeys", {})

‎utils/recover.py‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,9 @@ def calculate_keys(key_id, player_to_data, algo):
149149
else:
150150
raise RecoveryErrorUnknownAlgorithm(algo)
151151

152+
def _encode_algo_and_keyset(algo, keyset):
153+
return algo + ":" + str(keyset)
154+
152155
def restore_key_and_chaincode(zip_path, private_pem_path, passphrase, key_pass=None, mobile_key_pem_path = None, mobile_key_pass = None):
153156
privkeys = {}
154157
players_data = defaultdict(dict)
@@ -237,9 +240,9 @@ def restore_key_and_chaincode(zip_path, private_pem_path, passphrase, key_pass=N
237240
pub_from_metadata = signing_keys[key_id].public_key
238241
if pub_from_metadata != pubkey_str:
239242
print(f"Failed to recover {algo} key, expected public key is: {pub_from_metadata} calculated public key is: {pubkey_str}")
240-
privkeys[algo] = None
243+
privkeys[_encode_algo_and_keyset(algo, signing_keys[key_id].keyset_id)] = None
241244
else:
242-
privkeys[algo] = privkey, chain_code_for_this_key
245+
privkeys[_encode_algo_and_keyset(algo, signing_keys[key_id].keyset_id)] = privkey, chain_code_for_this_key, signing_keys[key_id].keyset_id
243246

244247
if len(privkeys) == 0:
245248
raise RecoveryErrorPublicKeyNoMatch()
@@ -281,3 +284,19 @@ def encode_extended_key(algo, key, chain_code, is_pub):
281284
extended_key += bytes(1)
282285
extended_key += key
283286
return encode_base58_checksum(extended_key)
287+
288+
def split_algo_and_keyset(algo: str):
289+
split_algo = algo.split(":")
290+
if len(split_algo) == 1:
291+
return split_algo[0], None
292+
elif len(split_algo) == 2:
293+
return split_algo[0], int(split_algo[1])
294+
else:
295+
raise ValueError("invalid algo encoding: '%s'" % algo)
296+
297+
def strip_keyset_from_results(privkeys: str):
298+
res = {}
299+
for algo, info in privkeys.items():
300+
res[split_algo_and_keyset(algo)[0]] = info
301+
return res
302+

0 commit comments

Comments
 (0)