#!/usr/bin/env python3
"""Generate / verify DVeProto test_vectors.json (pip install cryptography)."""
from __future__ import annotations

import json
import os
import struct
import sys
from pathlib import Path

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

HERE = Path(__file__).resolve().parent
VECTORS = HERE / "test_vectors.json"


def hkdf(ikm: bytes, info: bytes) -> bytes:
    return HKDF(algorithm=hashes.SHA256(), length=32, salt=b"", info=info).derive(ikm)


def main(write: bool = True) -> int:
    data = json.loads(VECTORS.read_text(encoding="utf-8"))
    ikm = bytes.fromhex(data["hkdf"]["ikm_hex"])
    c2s = hkdf(ikm, data["hkdf"]["info_c2s"].encode())
    s2c = hkdf(ikm, data["hkdf"]["info_s2c"].encode())
    data["hkdf"]["c2s_hex"] = c2s.hex()
    data["hkdf"]["s2c_hex"] = s2c.hex()

    token = bytes.fromhex(data["resume_hkdf"]["token_hex"])
    rc2s = hkdf(token, data["resume_hkdf"]["info_c2s"].encode())
    rs2c = hkdf(token, data["resume_hkdf"]["info_s2c"].encode())
    data["resume_hkdf"]["c2s_hex"] = rc2s.hex()
    data["resume_hkdf"]["s2c_hex"] = rs2c.hex()

    # --- 1.3 frame ---
    w13 = data["wire_1_3"]
    prefix = bytes.fromhex(w13["send_prefix_hex"])
    counter = int(w13["send_counter"])
    nonce = prefix + struct.pack(">Q", counter)
    aad = bytes([0x13, int(w13["ptype"]) & 0xFF]) + struct.pack(
        ">H", int(w13["stream_id"]) & 0xFFFF
    )
    assert aad.hex() == w13["aad_hex"], (aad.hex(), w13["aad_hex"])
    assert nonce.hex() == w13["nonce_hex"], (nonce.hex(), w13["nonce_hex"])
    pt = w13["plaintext_utf8"].encode("utf-8")
    ct = AESGCM(c2s).encrypt(nonce, pt, aad)
    w13["ciphertext_hex"] = ct.hex()

    # --- 1.4 frame ---
    w14 = data["wire_1_4"]
    prefix = bytes.fromhex(w14["send_prefix_hex"])
    counter = int(w14["send_counter"])
    pkt_seq = int(w14["pkt_seq"])
    nonce = prefix + struct.pack(">Q", counter)
    ptype = int(w14["ptype"]) & 0xFF
    sid = int(w14["stream_id"]) & 0xFFFF
    aad = bytes([0x14, ptype]) + struct.pack(">HI", sid, pkt_seq & 0xFFFFFFFF)
    assert aad.hex() == w14["aad_hex"], (aad.hex(), w14["aad_hex"])
    assert nonce.hex() == w14["nonce_hex"]
    pt = w14["plaintext_utf8"].encode("utf-8")
    ct = AESGCM(c2s).encrypt(nonce, pt, aad)
    frame = bytearray(20 + len(ct))
    frame[0] = 0x14
    frame[1] = ptype
    frame[2] = (sid >> 8) & 0xFF
    frame[3] = sid & 0xFF
    struct.pack_into(">I", frame, 4, pkt_seq)
    frame[8:20] = nonce
    frame[20:] = ct
    w14["frame_hex"] = bytes(frame).hex()

    if write:
        VECTORS.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8")
        print("wrote", VECTORS)
    print("c2s", data["hkdf"]["c2s_hex"][:16] + "…")
    print("1.4 frame len", len(frame))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(write="--check" not in sys.argv))