import { GetPublicKeyCommand, KMSClient } from "@aws-sdk/client-kms";
import { BitString, ObjectIdentifier, Sequence, verifySchema } from "asn1js";
import { getAddress, keccak256 } from "viem";

/**
 * ASN1 schema to parse the public key
 */
const EcdsaPubKey = new Sequence({
    name: "EcdsaPubKey",
    value: [
        new Sequence({
            name: "algo",
            value: [
                new ObjectIdentifier({ name: "a" }),
                new ObjectIdentifier({ name: "b" }),
            ],
        }),
        new BitString({ name: "pubKey" }),
    ],
});
/**
 * Get the ethereum address from the given KMS key id
 * @param keyId
 * @param client
 */
export async function getKmsAddress({
    keyId,
    client,
}: {
    keyId: string;
    client: KMSClient;
}) {
    const res = await client.send(
        new GetPublicKeyCommand({
            KeyId: keyId,
        })
    );
    if (!res.PublicKey) {
        throw new Error("Missing public key");
    }
    const publicKey = Buffer.from(res.PublicKey);

    // Ensure the signature match our expected format
    const decodedPubKey = verifySchema(publicKey, EcdsaPubKey);
    if (!decodedPubKey.verified) {
        throw new Error("Invalid public key");
    }
    // Get the raw public key, and then build our address buffer
    const rawPubKey = (
        decodedPubKey.result.valueBlock as any
    ).value[1].valueBlock.valueHexView.slice(1);
    const addressBuf = Buffer.from(keccak256(rawPubKey).slice(2), "hex");

    // Finally, map the address buffer to an ethereum address
    return getAddress(`0x${addressBuf.subarray(-20).toString("hex")}`);
}