discover_kms_root.py

#!/usr/bin/env python3
"""Recover KMS root address from a running CVM's signature chain."""

import sys
from dstack_sdk import DstackClient
from eth_keys import keys
from eth_utils import keccak

client = DstackClient()
info = client.info()
app_id = info.app_id
print(f"App ID: {app_id}")

result = client.get_key("/discover", "ethereum")
derived_key = bytes.fromhex(result.key.replace('0x', ''))[:32]
priv = keys.PrivateKey(derived_key)
derived_pubkey = priv.public_key.to_compressed_bytes()

app_sig = bytes.fromhex(result.signature_chain[0].replace('0x', ''))
kms_sig = bytes.fromhex(result.signature_chain[1].replace('0x', ''))

app_msg = f"ethereum:{derived_pubkey.hex()}"
app_msg_hash = keccak(text=app_msg)
app_pubkey = keys.Signature(app_sig).recover_public_key_from_msg_hash(app_msg_hash).to_compressed_bytes()

app_id_bytes20 = bytes.fromhex(app_id.replace('0x', ''))
kms_msg = b"dstack-kms-issued:" + app_id_bytes20 + app_pubkey
kms_msg_hash = keccak(kms_msg)
kms_signer = keys.Signature(kms_sig).recover_public_key_from_msg_hash(kms_msg_hash)

print(f"KMS root address: {kms_signer.to_checksum_address()}")
print(f"App pubkey: {app_pubkey.hex()}")
print(f"Derived pubkey: {derived_pubkey.hex()}")