#!/usr/bin/env python3
"""Reference vector generator for SCK AES-256-GCM wire format.

Independent of the RustCrypto suite so the Rust implementation
is checked against an external reference, not itself.

Wire format produced:
    [version:u8=0x01] [nonce:12] [ciphertext || tag(16)]

AAD = tenant_id.bytes (16, RFC 4122 big-endian)
    || config_uuid.bytes (16, RFC 4122 big-endian)
    || key_version.to_bytes(8, 'big', signed=True)     # i64 big-endian

Library: pyca/cryptography (OpenSSL backing) — AES-256-GCM.
"""

import uuid
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

def enc(sck: bytes, plaintext: bytes, nonce: bytes,
        tenant_id: uuid.UUID, config_uuid: uuid.UUID,
        key_version: int) -> bytes:
    assert len(sck) == 32
    assert len(nonce) == 12
    aad = tenant_id.bytes + config_uuid.bytes + key_version.to_bytes(8, "big", signed=True)
    assert len(aad) == 40
    aead = AESGCM(sck)
    ct_and_tag = aead.encrypt(nonce, plaintext, aad)  # tag appended
    return bytes([0x01]) + nonce + ct_and_tag

def emit(label, **kwargs):
    wire = enc(**kwargs)
    print(f"# {label}")
    print(f"#   sck         = {kwargs['sck'].hex()}")
    print(f"#   nonce       = {kwargs['nonce'].hex()}")
    print(f"#   tenant_id   = {kwargs['tenant_id']}")
    print(f"#   config_uuid = {kwargs['config_uuid']}")
    print(f"#   key_version = {kwargs['key_version']}")
    print(f"#   plaintext   = {kwargs['plaintext']!r}")
    print(f"#   plaintext hex = {kwargs['plaintext'].hex()}")
    print(f"#   wire len    = {len(wire)} bytes")
    print(f"#   wire hex    =")
    hx = wire.hex()
    for i in range(0, len(hx), 64):
        print(f"#       {hx[i:i+64]}")
    print()

if __name__ == "__main__":
    print(f"# generated by gen_sck.py (cryptography {__import__('cryptography').__version__})\n")

    emit("vector 1 — simple JSON, 38-byte plaintext",
         sck=bytes(range(0x00, 0x20)),
         plaintext=br'{"realm":"llm","impl":"x","config":{}}',
         nonce=bytes(range(0x10, 0x1c)),
         tenant_id=uuid.UUID("11111111-1111-1111-1111-111111111111"),
         config_uuid=uuid.UUID("22222222-2222-2222-2222-222222222222"),
         key_version=1)

    emit("vector 2 — Unicode JSON (Japanese characters in a value)",
         sck=bytes(range(0x20, 0x40)),
         plaintext='{"display_name":"テスト"}'.encode("utf-8"),
         nonce=bytes(range(0x30, 0x3c)),
         tenant_id=uuid.UUID("33333333-3333-3333-3333-333333333333"),
         config_uuid=uuid.UUID("44444444-4444-4444-4444-444444444444"),
         key_version=7)

    emit("vector 3 — edge case, shortest conceivable JSON blob",
         sck=bytes(range(0x40, 0x60)),
         plaintext=b"{}",
         nonce=bytes(range(0x50, 0x5c)),
         tenant_id=uuid.UUID("55555555-5555-5555-5555-555555555555"),
         config_uuid=uuid.UUID("66666666-6666-6666-6666-666666666666"),
         key_version=42)
