#!/usr/bin/env python3
"""
LDN Data Frame Key Derivation for Mario Maker 2
Requires: pip install pycryptodome scapy
"""

import hashlib
from Crypto.Cipher import AES

# ── Nintendo key constants ────────────────────────────────────────────────────
MASTER_KEY_00            = bytes.fromhex("c2caaff089b9aed55694876055271c7d")
AES_KEK_GENERATION_SRC   = bytes.fromhex("4d870986c45d20722fba1053da92e8a9")
AES_KEY_GENERATION_SRC   = bytes.fromhex("89615ee05c31b6805fe58f3da24f7aa8")
DATA_INPUT_KEY           = bytes.fromhex("f1e7018419a84f711da714c2cf919c9c")
NINTENDO_OUI             = bytes([0x00, 0x22, 0xAA])

# ── Game-specific ─────────────────────────────────────────────────────────────
PASSWORD = b"LunchPack2DefaultPhrase"

# ── AES-ECB helper ───────────────────────────────────────────────────────────
def aes_ecb_decrypt(key: bytes, data: bytes) -> bytes:
    return AES.new(key, AES.MODE_ECB).decrypt(data)

# ── Step-by-step key derivation ──────────────────────────────────────────────
def derive_data_key(server_random: bytes, password: bytes) -> bytes:
    assert len(server_random) == 16, "server_random must be 16 bytes"

    # Step 1: kek_key
    kek_key = aes_ecb_decrypt(MASTER_KEY_00, AES_KEK_GENERATION_SRC)
    print(f"[1] kek_key:         {kek_key.hex()}")

    # Step 2: access_key (decrypt the data input key with kek_key)
    access_key = aes_ecb_decrypt(kek_key, DATA_INPUT_KEY)
    print(f"[2] access_key:      {access_key.hex()}")

    # Step 3: generation_key
    generation_key = aes_ecb_decrypt(access_key, AES_KEY_GENERATION_SRC)
    print(f"[3] generation_key:  {generation_key.hex()}")

    # Step 4: hash input buffer = server_random || password
    buf = server_random + password
    hash_input = hashlib.sha256(buf).digest()[:16]
    print(f"[4] SHA256(buf)[0:16]: {hash_input.hex()}")

    # Step 4: final data frame key
    data_key = aes_ecb_decrypt(generation_key, hash_input)
    print(f"[★] data_frame_key:  {data_key.hex()}")

    return data_key

# ── Main ─────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import sys
    pcap_path = sys.argv[1] if len(sys.argv) > 1 else "./capture.pcapng"

    server_random = bytes.fromhex("aaec2d390000000005180000796265e7")

    data_key = derive_data_key(bytes.fromhex("d33f9c317a7f2b77a27b68e40f596f89"), PASSWORD)

    print(f"tk: {data_key.hex()}")