#!/usr/bin/env -S uv run --quiet --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["noiseprotocol>=0.3.1", "pynacl>=1.5"]
# ///
"""Generate test/vectors.json for the JS implementation.

The JS crypto in web/js must agree byte for byte with the reference stack
(hashlib/hmac, PyNaCl, noiseprotocol), so the ground truth is generated here
with fixed inputs and consumed by test/vectors.test.mjs. Everything is
derived from SHA-256 of fixed labels, so the file is reproducible.
"""

from __future__ import annotations

import hashlib
import hmac
import json
import struct
from pathlib import Path

import nacl.bindings as sodium
from noise.connection import Keypair, NoiseConnection

LABEL_SEAL, LABEL_MSG, LABEL_ID = b"smolmail/1 seal", b"smolmail/1 msg", b"smolmail/1 id"
LABEL_IDENTITY, LABEL_ACCEPT, LABEL_MAC = b"smolmail/1 identity", b"smolmail/1 accept", b"smolmail/1 mac"
TIME = 1730000000  # fixed so the envelope vector is reproducible


def seed(label: str) -> bytes:
    return hashlib.sha256(label.encode()).digest()


def hkdf_sha256(ikm: bytes, salt: bytes, info: bytes, length: int) -> bytes:
    prk = hmac.new(salt, ikm, hashlib.sha256).digest()
    out, block, counter = b"", b"", 1
    while len(out) < length:
        block = hmac.new(prk, block + info + bytes([counter]), hashlib.sha256).digest()
        out, counter = out + block, counter + 1
    return out[:length]


def main() -> int:
    out: dict = {}

    # --- hashes over empty, one-block and multi-block inputs
    out["sha256"] = [
        {"in": m.hex(), "out": hashlib.sha256(m).hexdigest()}
        for m in (b"", b"abc", bytes(range(64)), b"smolmail" * 40)
    ]
    out["sha512"] = [
        {"in": m.hex(), "out": hashlib.sha512(m).hexdigest()}
        for m in (b"", b"abc", bytes(range(64)))
    ]

    # --- HMAC/HKDF, including RFC 5869 test case 1
    out["hkdf"] = []
    for ikm, salt, info, length, name in (
        (bytes.fromhex("0b" * 22), bytes.fromhex("000102030405060708090a0b0c"),
         bytes.fromhex("f0f1f2f3f4f5f6f7f8f9"), 42, "rfc5869-1"),
        (b"agreement shared", b"epk" + b"recipient", b"smolmail/1 seal", 32, "seal-shaped"),
    ):
        out["hkdf"].append({"name": name, "ikm": ikm.hex(), "salt": salt.hex(),
                            "info": info.hex(), "len": length,
                            "out": hkdf_sha256(ikm, salt, info, length).hex()})

    # --- ChaCha20-Poly1305 (RFC 8439 IETF AEAD) over varied shapes
    key, nonce = seed("aead key"), seed("aead nonce")[:12]
    aad = bytes.fromhex("50515253c0c1c2c3c4c5c6c7")
    out["aead"] = []
    for plaintext, use_aad, name in (
        (b"", False, "empty"), (b"hello", True, "short"),
        (b"Ladies and Gentlemen of the class of '99: if I could offer you "
         b"only one tip for the future, sunscreen would be it.", True, "multiline"),
    ):
        sealed = sodium.crypto_aead_chacha20poly1305_ietf_encrypt(
            plaintext, aad if use_aad else b"", nonce, key)
        out["aead"].append({"name": name, "key": key.hex(), "nonce": nonce.hex(),
                            "aad": (aad if use_aad else b"").hex(), "plaintext": plaintext.hex(),
                            "sealed": sealed.hex()})

    # --- X25519: base multiplication and agreement, including the §2 checks
    out["x25519"] = []
    for name, priv in (("alice", seed("x25519 alice")), ("bob", seed("x25519 bob"))):
        pub = sodium.crypto_scalarmult_base(priv)
        out["x25519"].append({"name": name, "priv": priv.hex(), "pub": pub.hex()})
    out["x25519"].append({"name": "agree",
                          "shared": sodium.crypto_scalarmult(
                              seed("x25519 alice"),
                              sodium.crypto_scalarmult_base(seed("x25519 bob"))).hex()})
    for bad in (b"\x00" * 32, bytes.fromhex("0100" + "00" * 30)):
        try:
            sodium.crypto_scalarmult(seed("x25519 alice"), bad)
            raised = False
        except Exception:
            raised = True
        out["x25519"].append({"name": f"low-order-{bad[0]}", "peer": bad.hex(),
                              "peer_rejected": True, "raised": raised})
        assert raised, "PyNaCl must reject this low-order point for the vector to mean anything"

    # --- Ed25519: keypair, signature, verification, and a tampered case
    out["ed25519"] = []
    for name, s in (("vector-seed-a", seed("ed25519 a")), ("vector-seed-b", seed("ed25519 b"))):
        pub, _ = sodium.crypto_sign_seed_keypair(s)
        for message in (b"", b"smolmail/1 auth" + bytes(32), bytes(range(64))):
            sig = sodium.crypto_sign(message, s + pub)[:64]
            out["ed25519"].append({"name": name, "seed": s.hex(), "pub": pub.hex(),
                                   "message": message.hex(), "signature": sig.hex(),
                                   "valid": True})
        bad = bytes([sig[0] ^ 1]) + sig[1:]
        out["ed25519"].append({"name": name, "seed": s.hex(), "pub": pub.hex(),
                               "message": message.hex(), "signature": bad.hex(),
                               "valid": False})

    # --- §2 conversions between the identity key and X25519
    out["ed_to_x25519"] = []
    for name, s in (("vector-seed-a", seed("ed25519 a")), ("vector-seed-b", seed("ed25519 b"))):
        pub, sk64 = sodium.crypto_sign_seed_keypair(s)
        out["ed_to_x25519"].append({
            "name": name, "seed": s.hex(),
            "x_priv": sodium.crypto_sign_ed25519_sk_to_curve25519(sk64).hex(),
            "x_pub": sodium.crypto_sign_ed25519_pk_to_curve25519(pub).hex()})

    # --- §5 envelope sealed with a fixed ephemeral, mirroring smolmail.py seal()
    sender_seed, recipient_seed, esk = seed("envelope sender"), seed("envelope recipient"), seed("envelope esk")
    sender_pub, _ = sodium.crypto_sign_seed_keypair(sender_seed)
    recipient_pub, _ = sodium.crypto_sign_seed_keypair(recipient_seed)
    epk = sodium.crypto_scalarmult_base(esk)
    shared = sodium.crypto_scalarmult(esk, sodium.crypto_sign_ed25519_pk_to_curve25519(recipient_pub))
    key = hkdf_sha256(shared, epk + recipient_pub, LABEL_SEAL, 32)
    body = b"---\nSubject: vector\n---\nhello bob\n"
    sig = sodium.crypto_sign(LABEL_MSG + recipient_pub + epk
                             + bytes([1]) + sender_pub + struct.pack(">qI", TIME, len(body)) + body,
                             sender_seed + sender_pub)[:64]
    plaintext = (bytes([1]) + sender_pub + struct.pack(">qI", TIME, len(body))
                 + body + sig)
    aad = b"SMOL" + bytes([1]) + recipient_pub + epk
    envelope = aad + sodium.crypto_aead_chacha20poly1305_ietf_encrypt(plaintext, aad, bytes(12), key)
    out["envelope"] = {
        "sender_seed": sender_seed.hex(), "recipient_seed": recipient_seed.hex(),
        "esk": esk.hex(), "body": body.hex(), "time": TIME,
        "envelope": envelope.hex(),
        "id": hashlib.sha256(LABEL_ID + envelope).digest().hex(),
        "unpadded_plaintext_len": len(plaintext)}

    # --- §2/§5.8: master-derived rotation seeds and accept tokens
    master = seed("master")
    correspondent = seed("correspondent identity")
    accept_key = hkdf_sha256(master, b"", LABEL_ACCEPT, 32)
    token = hmac.new(accept_key, correspondent, hashlib.sha256).digest()
    mid = seed("token message id")
    out["tokens"] = {
        "master": master.hex(),
        "identity_seed_0": hkdf_sha256(master, b"", LABEL_IDENTITY + struct.pack(">I", 0), 32).hex(),
        "identity_seed_1": hkdf_sha256(master, b"", LABEL_IDENTITY + struct.pack(">I", 1), 32).hex(),
        "accept_key": accept_key.hex(),
        "correspondent": correspondent.hex(),
        "token": token.hex(),
        "message_id": mid.hex(),
        "accept_mac": hmac.new(token, LABEL_MAC + mid, hashlib.sha256).digest().hex(),
    }

    # --- Noise NX transcript with every key fixed
    i_e, r_e, s_priv = seed("nx initiator eph"), seed("nx responder eph"), seed("nx server static")
    s_pub = sodium.crypto_scalarmult_base(s_priv)
    n1 = NoiseConnection.from_name(b"Noise_NX_25519_ChaChaPoly_SHA256")
    n1.set_as_initiator()
    n1.set_prologue(b"smolmail/1")
    n1.set_keypair_from_private_bytes(Keypair.EPHEMERAL, i_e)
    n1.start_handshake()
    m1 = n1.write_message()

    n2 = NoiseConnection.from_name(b"Noise_NX_25519_ChaChaPoly_SHA256")
    n2.set_as_responder()
    n2.set_prologue(b"smolmail/1")
    n2.set_keypair_from_private_bytes(Keypair.STATIC, s_priv)
    n2.set_keypair_from_private_bytes(Keypair.EPHEMERAL, r_e)
    n2.start_handshake()
    n2.read_message(m1)
    m2 = n2.write_message()
    n1.read_message(m2)
    assert n1.handshake_finished and n2.handshake_finished

    out["noise"] = {
        "initiator_eph_priv": i_e.hex(), "responder_eph_priv": r_e.hex(),
        "server_static_priv": s_priv.hex(), "server_static_pub": s_pub.hex(),
        "message1": m1.hex(), "message2": m2.hex(),
        "handshake_hash": n1.get_handshake_hash().hex(),
        "assert_hash_equal": n1.get_handshake_hash() == n2.get_handshake_hash(),
        "initiator_frames": [{"plaintext": p.hex(), "sealed": n1.encrypt(p).hex()}
                             for p in (b"ping", b"second frame to test the counter")],
        "responder_frames": [{"plaintext": p.hex(), "sealed": n2.encrypt(p).hex()}
                             for p in (b"pong", b"x" * 300)],
    }

    path = Path(__file__).parent / "vectors.json"
    path.write_text(json.dumps(out, indent=2) + "\n")
    print(f"wrote {path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())