#!/usr/bin/env python3
# pip install pycryptodome
import hashlib
from Crypto.Cipher import AES

# --- Constants (public, from prod.keys) ---
AES_KEK_GENERATION_SOURCE = bytes.fromhex("4d870986c45d20722fba1053da92e8a9")
AES_KEY_GENERATION_SOURCE  = bytes.fromhex("89615ee05c31b6805fe58f3da24f7aa8")

# LDN data frame input key (from kinnay's wiki)
LDN_INPUT_KEY = bytes.fromhex("f1e7018419a84f711da714c2cf919c9c")

def aes_unwrap(wrapped_key: bytes, wrap_key: bytes) -> bytes:
    """Single AES-128 block decryption (ECB)."""
    return AES.new(wrap_key, AES.MODE_ECB).decrypt(wrapped_key)

def derive_ldn_key(master_key_00: bytes, server_random: bytes, passphrase: str) -> bytes:
    """
    Derives the AES-CTR key used to encrypt LDN data frames.
    Steps follow the LDN Protocol wiki exactly.
    """
    # Step 1: decrypt aes_kek_generation_source with master_key_00
    kek = aes_unwrap(AES_KEK_GENERATION_SOURCE, master_key_00)

    # Step 2: decrypt the LDN input key with the result
    kek2 = aes_unwrap(LDN_INPUT_KEY, kek)

    # Step 3: decrypt aes_key_generation_source with the result
    kek3 = aes_unwrap(AES_KEY_GENERATION_SOURCE, kek2)

    # Step 4: build the input buffer = server_random + passphrase bytes
    passphrase_bytes = passphrase.encode("utf-8")
    input_buffer = server_random + passphrase_bytes

    # Hash it, take first 16 bytes, decrypt with kek3
    hash_bytes = hashlib.sha256(input_buffer).digest()[:16]
    data_key = aes_unwrap(hash_bytes, kek3)

    return data_key


if __name__ == "__main__":
    # --- Fill these in ---
    # Your master_key_00 from prod.keys
    master_key_00 = bytes.fromhex("c2caaff089b9aed55694876055271c7d")

    # 16-byte server random from the LDN advertisement frame in Wireshark
    server_random = bytes.fromhex("YOUR_SERVER_RANDOM_HERE")

    passphrase = "LunchPack2DefaultPhrase"

    key = derive_ldn_key(master_key_00, server_random, passphrase)
    print(f"Derived AES-CTR key: {key.hex()}")