#!/usr/bin/env python3
# Регенерация tests/ayla/kdf_vectors.hpp. Эталон — APK AylaEncryption
# (совпадает с tools/probe_reference.py). Зависимость: pycryptodome.
# Запуск: python scripts/gen_kdf_vectors.py > tests/ayla/kdf_vectors.hpp
import base64, hmac, json, random, string
from Crypto.Cipher import AES

random.seed(20260922)

def rand_token(n=16, alpha=string.ascii_letters + string.digits):
    return "".join(random.choice(alpha) for _ in range(n))

def cstr(s):
    return json.dumps(s)

vectors = []
for i in range(4):
    lanip_key = base64.b64encode(random.randbytes(15 + i)).decode()
    rnd1 = rand_token(16, string.ascii_letters + string.digits + "+/")
    rnd2 = rand_token(16)
    t1 = random.randrange(10**12, 10**15)
    t2 = random.randrange(10**12, 10**15)
    k = lanip_key.encode(); b1, b2 = rnd1.encode(), rnd2.encode()
    s1, s2 = str(t1).encode(), str(t2).encode()

    def m(msg, suf):
        msg = msg + bytes([suf])
        return hmac.digest(k, hmac.digest(k, msg, "sha256") + msg, "sha256")

    A, D = b1 + b2 + s1 + s2, b2 + b1 + s2 + s1
    app_sign, dev_sign = m(A, 0x30), m(D, 0x30)
    e = AES.new(m(A, 0x31), AES.MODE_CBC, m(A, 0x32)[:16])
    msgs = [
        '{"seq_no":0,"data":{}}',
        '{"seq_no":1,"data":{"cmds":[{"cmd":{"method":"GET","resource":"property.json?name=fan_speed","uri":"/local_lan/property/datapoint.json","data":"","cmd_id":3}}]}}',
        '{"seq_no":2,"data":{"properties":[{"property":{"base_type":"integer","name":"fan_speed","value":3,"id":"ab12cd34"}}]}}',
    ]
    encs = []
    for mtext in msgs:
        raw = mtext.encode()
        n = ((len(raw) + 1 + 15) // 16) * 16          # Java-паддинг (>=1 NUL)
        encs.append(base64.b64encode(e.encrypt(raw.ljust(n, b"\x00"))).decode())
    legacy_msg = '{"seq_no":3,"data":{"pad":"XX"}}'   # 32 байта, без NUL
    assert len(legacy_msg) % 16 == 0 and legacy_msg[-1] != "\x00"
    encs.append(base64.b64encode(e.encrypt(legacy_msg.encode())).decode())
    msgs.append(legacy_msg)

    vectors.append(dict(
        lanip_key=lanip_key, random_1=rnd1, random_2=rnd2,
        time_1=t1, time_2=t2, plaintexts=msgs, enc_base64=encs,
        app_sign_msg0=base64.b64encode(hmac.digest(app_sign, msgs[0].encode(), "sha256")).decode(),
        dev_sign_msg1=base64.b64encode(hmac.digest(dev_sign, msgs[1].encode(), "sha256")).decode()))

lines = []
lines.append("// СГЕНЕРИРОВАНО scripts/gen_kdf_vectors.py (эталон = APK AylaEncryption).")
lines.append("// Не редактировать руками; регенерация:")
lines.append("//   python scripts/gen_kdf_vectors.py > tests/ayla/kdf_vectors.hpp")
lines.append("#pragma once")
lines.append("#include <cstddef>")
lines.append("")
lines.append("namespace fgl::ayla::test {")
lines.append("")
lines.append("struct KdfVector {")
lines.append("  const char* lanip_key;")
lines.append("  const char* random_1;")
lines.append("  const char* random_2;")
lines.append("  long long time_1;")
lines.append("  long long time_2;")
lines.append("  const char* plaintexts[4];")
lines.append("  const char* enc_base64[4];   // 0..2: Java-паддинг; 3: legacy без NUL")
lines.append("  const char* app_sign_msg0_b64;")
lines.append("  const char* dev_sign_msg1_b64;")
lines.append("};")
lines.append("")
lines.append("inline constexpr KdfVector kKdfVectors[] = {")
for v in vectors:
    pt = ", ".join(cstr(m) for m in v["plaintexts"])
    en = ", ".join(cstr(e) for e in v["enc_base64"])
    lines.append(f"  {{{cstr(v['lanip_key'])}, {cstr(v['random_1'])}, {cstr(v['random_2'])}, "
                 f"{v['time_1']}LL, {v['time_2']}LL, {{{pt}}}, {{{en}}}, "
                 f"{cstr(v['app_sign_msg0'])}, {cstr(v['dev_sign_msg1'])}}},")
lines.append("};")
lines.append("inline constexpr size_t kKdfVectorCount = sizeof(kKdfVectors) / sizeof(kKdfVectors[0]);")
lines.append("")
lines.append("}  // namespace fgl::ayla::test")
print("\n".join(lines))