"""UdbAESUtil 参考实现与设备实测对拍。

2026-08-25 Frida 直调差分结论（scripts/hook_aes_diff.py + evidence/aes_diff.json）:
  - UdbAESUtil::Cipher  == 标准 AES-128 加密 (FIPS-197 KAT 逐位一致)
  - UdbAESUtil::KeyExpansion == 标准展开，轮密钥按「列序/转置」存储
  - C1(uchar* key) 只读 key 前 16 字节
即: 不存在自定义 AES；任何"解密对不上"都是密钥来源/调用方的问题。
"""
from __future__ import annotations

from Crypto.Cipher import AES

MASTER_KEY_STR = b"HuyaUdb1928374650qwertyuiop"

_SBOX = bytes.fromhex(
    '637c777bf26b6fc53001672bfed7ab76ca82c97dfa5947f0add4a2af9ca472c0'
    'b7fd9326363ff7cc34a5e5f171d8311504c723c31896059a071280e2eb27b275'
    '09832c1a1b6e5aa0523bd6b329e32f8453d100ed20fcb15b6acbbe394a4c58cf'
    'd0efaafb434d338545f9027f503c9fa851a3408f929d38f5bcb6da2110fff3d2'
    'cd0c13ec5f974417c4a77e3d645d197360814fdc222a908846eeb814de5e0bdb'
    'e0323a0a4906245cc2d3ac629195e479e7c8376d8dd54ea96c56f4ea657aae08'
    'ba78252e1ca6b4c6e8dd741f4bbd8b8a703eb5664803f60e613557b986c11d9e'
    'e1f8981169d98e949b1e87e9ce5528df8ca1890dbfe6426841992d0fb054bb16')
_INV = bytes.fromhex(
    '52096ad53036a538bf40a39e81f3d7fb733317225e060830ca22c1e0941053ff'
    '5013071554c6935c9b6b478e49af25ebd6d13ca8d1b972e183a15b8e35997231'
    'd18def6e281286be5532e377b905b09393c428c0aab8858fddec191bdbaff54c'
    'bb617a475e8db4f89d389128723abbd6142cbbb8da24014361a9cf83d0d90c0b'
    '93ae82fdf3a50165c19fd02fa1ca67094c40b5f9bbc9ecd454f75bc46caa5e58'
    'f4fc6d7893fa59957c04c0a3a8ef7a3493163b536fceaac2d940e70fa65048e9'
    '4bccd5d83760a68309a2ded466b75c8976ee06753c70336a124b2042b8db9a8f'
    '58eec4f73d116f33766f2cc10ed1233b84f21f2679120569ad47eb4e30665ff1')


def _xtime(a: int) -> int:
    a <<= 1
    return (a ^ 0x1B) & 0xFF if a & 0x100 else a


def expand_schedule_transposed(key16: bytes) -> bytes:
    """标准 AES-128 轮密钥，按 UdbAESUtil 的存储序输出（176B）。

    设备布局: rk[round] 为 16B，按 state 列序存放（rk0..3 = key 的第0列）。
    """
    assert len(key16) == 16
    words = [list(key16[4 * i:4 * i + 4]) for i in range(4)]
    rcon = 1
    for i in range(4, 44):
        t = list(words[i - 1])
        if i % 4 == 0:
            t = t[1:] + t[:1]
            t = [_SBOX[b] for b in t]
            t[0] ^= rcon
            rcon = _xtime(rcon)
        words.append([words[i - 4][j] ^ t[j] for j in range(4)])
    # 设备实测布局: 每轮16B按字表转置存放 dev[r*4+c] = words[rnd*4+c][r]
    out = bytearray()
    for rnd in range(11):
        blk = bytearray(16)
        for c in range(4):
            for r in range(4):
                blk[r * 4 + c] = words[rnd * 4 + c][r]
        out += blk
    return bytes(out)


def encrypt_block(key16: bytes, block16: bytes) -> bytes:
    return AES.new(key16, AES.MODE_ECB).encrypt(block16)


def decrypt_block(key16: bytes, block16: bytes) -> bytes:
    return AES.new(key16, AES.MODE_ECB).decrypt(block16)


def zero_pad(data: bytes) -> bytes:
    """UdbAESUtil::encrypt 的填充：补 '\\0' 到 16 的倍数（不追加整块）。"""
    if len(data) % 16 == 0:
        return data
    return data + b"\x00" * (16 - len(data) % 16)


def udb_encrypt(key16: bytes, plain: bytes) -> bytes:
    p = zero_pad(plain)
    return AES.new(key16, AES.MODE_ECB).encrypt(p)


def udb_decrypt(key16: bytes, ct: bytes) -> bytes:
    return AES.new(key16, AES.MODE_ECB).decrypt(ct).rstrip(b"\x00")


def selftest() -> None:
    # FIPS-197
    kat_k = bytes(range(16))
    kat_p = bytes.fromhex("00112233445566778899aabbccddeeff")
    assert encrypt_block(kat_k, kat_p).hex() == "69c4e0d86a7b0430d8cdb78070b4c55a"
    # 与设备 dump 的轮密钥对拍
    import json
    from pathlib import Path
    j = Path(__file__).resolve().parent.parent / "evidence" / "aes_diff.json"
    if j.exists():
        dev = json.loads(j.read_text())["sched_kat"]["rk"]
        assert expand_schedule_transposed(kat_k).hex() == dev, "schedule mismatch!"
        dev_ct = json.loads(j.read_text())["cipher_huyaudb_first16"]["after"]
        # harness 该样本实际喂入的 key = "4875796155646231"(hex串被二次编码, 见脚本注释)
        ref = encrypt_block(b"4875796155646231", b"\x00" * 16).hex()
        assert dev_ct == ref, f"huyaudb mismatch {dev_ct} vs {ref}"
        # 主密钥前16字节的标准指纹(供真实样本比对)
        ref_master = encrypt_block(MASTER_KEY_STR[:16], b"\x00" * 16).hex()
        print("master first16 enc(zeros) =", ref_master)
    print("udb_aes selftest OK")


if __name__ == "__main__":
    selftest()