#!/usr/bin/env python3
# Author:  cirakg & claude ai
# Summary: Generate test/data/Vectors.mc: official crypto/OTP test vectors, cross-checked, plus seeded extra cases.
"""Generate the Monkey C test vectors in test/data/Vectors.mc.

Sources, and how each is double-checked before it is written:

* Published vectors (FIPS 180-4 examples, RFC 2202, RFC 4231, RFC 7914 §11,
  RFC 5869, RFC 4648 §10, RFC 4226 Appendix D, RFC 6238 Appendix B, NIST SP
  800-38A F.2.1/F.2.5) are written into this file from the documents. Python's
  standard library (hashlib, hmac, base64) recomputes each one, and the
  generator stops if the two disagree. AES isn't in the standard library, so
  the NIST AES vectors are checked with `openssl enc` instead.
* Extra cases (random messages around the SHA padding boundaries, random OTP
  secrets, digit counts, periods, times around 2038, Steam codes) come from a
  fixed seed and are computed with the standard library only, so they are
  reproducible and independent of the Monkey C code under test.

Usage:  gen_vectors.py [--out test/data/Vectors.mc]
The output is committed; rerun after changing this file. All data is fake.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import random
import struct
import subprocess
from collections.abc import Sequence
from pathlib import Path

SEED = 20261018
ALGS = {"SHA1": 0, "SHA256": 1, "SHA512": 2}
STEAM_ALPHABET = "23456789BCDFGHJKMNPQRTVWXY"


def h(data: bytes) -> str:
    return data.hex()


# --------------------------------------------------------------------- hashes

FIPS_MESSAGES = [
    b"",
    b"abc",
    b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq",
    b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu",
]
# FIPS 180-4 example digests of "abc", as printed in the standard's examples.
FIPS_ABC = {
    "SHA1": "a9993e364706816aba3e25717850c26c9cd0d89d",
    "SHA256": "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
    "SHA512": "ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a"
    "2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f",
}


def hash_vectors(rng: random.Random) -> list[tuple[int, str, str]]:
    for alg, expected in FIPS_ABC.items():
        assert hashlib.new(alg.lower(), b"abc").hexdigest() == expected, alg
    messages = list(FIPS_MESSAGES)
    # Lengths around the padding boundaries of 64-byte (SHA-1/256) and
    # 128-byte (SHA-512) blocks, where implementations usually break.
    for length in (55, 56, 63, 64, 65, 111, 112, 119, 120, 127, 128, 129, 200, 255, 256):
        messages.append(bytes(rng.randrange(256) for _ in range(length)))
    return [(ALGS[a], h(m), hashlib.new(a.lower(), m).hexdigest()) for a in ALGS for m in messages]


# ----------------------------------------------------------------------- HMAC

KEY_20 = b"\x0b" * 20
KEY_AA131 = b"\xaa" * 131
RFC4231_CASES = [  # (key, data); case 5's output is truncated to 128 bits in the RFC
    (KEY_20, b"Hi There"),
    (b"Jefe", b"what do ya want for nothing?"),
    (b"\xaa" * 20, b"\xdd" * 50),
    (bytes(range(1, 26)), b"\xcd" * 50),
    (b"\x0c" * 20, b"Test With Truncation"),
    (KEY_AA131, b"Test Using Larger Than Block-Size Key - Hash Key First"),
    (
        KEY_AA131,
        b"This is a test using a larger than block-size key and a larger than block-size data."
        b" The key needs to be hashed before being used by the HMAC algorithm.",
    ),
]
RFC2202_SHA1_CASES = [
    (KEY_20, b"Hi There"),
    (b"Jefe", b"what do ya want for nothing?"),
    (b"\xaa" * 20, b"\xdd" * 50),
    (bytes(range(1, 26)), b"\xcd" * 50),
    (b"\x0c" * 20, b"Test With Truncation"),
    (b"\xaa" * 80, b"Test Using Larger Than Block-Size Key - Hash Key First"),
    (b"\xaa" * 80, b"Test Using Larger Than Block-Size Key and Larger Than One Block-Size Data"),
]
# Spot checks copied from the RFCs.
RFC_HMAC_CHECKS = [
    ("SHA1", RFC2202_SHA1_CASES[0], "b617318655057264e28bc0b6fb378c8ef146be00"),
    ("SHA256", RFC4231_CASES[0], "b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7"),
    (
        "SHA512",
        RFC4231_CASES[0],
        "87aa7cdea5ef619d4ff0b4241a1d6cb02379f4e2ce4ec2787ad0b30545e17cde"
        "daa833b7d6b8a702038b274eaea3f4e4be9d914eeb61f1702e696c203a126854",
    ),
    ("SHA256", RFC4231_CASES[5], "60e431591ee0b67f0d8a26aacbf5b77f8e0bc6213728c5140546040f0ee37f54"),
]


def hmac_vectors() -> list[tuple[int, str, str, str]]:
    for alg, (key, data), expected in RFC_HMAC_CHECKS:
        assert hmac.new(key, data, alg.lower()).hexdigest() == expected, alg
    out = [(ALGS["SHA1"], h(k), h(d), hmac.new(k, d, "sha1").hexdigest()) for k, d in RFC2202_SHA1_CASES]
    out += [
        (ALGS[a], h(k), h(d), hmac.new(k, d, a.lower()).hexdigest())
        for a in ("SHA256", "SHA512")
        for k, d in RFC4231_CASES
    ]
    return out


# --------------------------------------------------------------- PBKDF2, HKDF

RFC7914_PBKDF2 = (
    b"passwd",
    b"salt",
    1,
    64,
    "55ac046e56e3089fec1691c22544b605f94185216dde0465e68b9d57c20dacbc"
    "49ca9cccf179b645991664b39d77ef317c71b845b1e30bd509112041d3a19783",
)


def pbkdf2_vectors() -> list[tuple[str, str, int, int, str]]:
    p, s, c, n, expected = RFC7914_PBKDF2
    assert hashlib.pbkdf2_hmac("sha256", p, s, c, n).hex() == expected
    cases = [
        (p, s, c, n),
        (b"password", b"salt", 2, 32),
        (b"password", b"salt", 100, 32),
        (b"123456", b"\x00" * 16, 1000, 32),
    ]
    cases.append((b"passwordPASSWORDpassword", b"saltSALTsaltSALTsaltSALTsaltSALTsalt", 300, 40))
    return [(h(p), h(s), c, n, hashlib.pbkdf2_hmac("sha256", p, s, c, n).hex()) for p, s, c, n in cases]


def hkdf(ikm: bytes, salt: bytes, info: bytes, length: int) -> bytes:
    prk = hmac.new(salt or b"\x00" * 32, ikm, "sha256").digest()
    okm, block = b"", b""
    for counter in range(1, -(-length // 32) + 1):
        block = hmac.new(prk, block + info + bytes([counter]), "sha256").digest()
        okm += block
    return okm[:length]


RFC5869_CASES = [
    (KEY_20 + b"\x0b" * 2, bytes(range(13)), bytes(range(0xF0, 0xFA)), 42),
    (bytes(range(0x50)), bytes(range(0x60, 0xB0)), bytes(range(0xB0, 0x100)), 82),
    (KEY_20 + b"\x0b" * 2, b"", b"", 42),
]
RFC5869_OKM = {
    0: "3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865",
    2: "8da4e775a563c18f715f802a063c5a31b8a11f5c5ee1879ec3454e5f3c738d2d9d201395faa4b61a96c8",
}


def hkdf_vectors() -> list[tuple[str, str, str, int, str]]:
    for index, expected in RFC5869_OKM.items():
        ikm, salt, info, n = RFC5869_CASES[index]
        assert hkdf(ikm, salt, info, n).hex() == expected, index
    return [(h(i), h(s), h(f), n, hkdf(i, s, f, n).hex()) for i, s, f, n in RFC5869_CASES]


# ------------------------------------------------------------------------ AES

NIST_IV = "000102030405060708090a0b0c0d0e0f"
NIST_PT = (
    "6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e51"
    "30c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710"
)
NIST_CBC = [  # SP 800-38A F.2.1 (AES-128) and F.2.5 (AES-256)
    (
        "2b7e151628aed2a6abf7158809cf4f3c",
        "7649abac8119b246cee98e9b12e9197d5086cb9b507219ee95db113a917678b2"
        "73bed6b8e3c1743b7116e69e222295163ff1caa1681fac09120eca307586e1a7",
    ),
    (
        "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
        "f58c4c04d6e5f1ba779eabfb5f7bfbd69cfc4e967edb808d679f777bc6702c7d"
        "39f23369a9d9bacfa530e26304231461b2eb05e2c39be9fcda6c19078c6a9d1b",
    ),
]


def openssl_cbc(key: str, iv: str, plaintext: str) -> str:
    cipher = "aes-128-cbc" if len(key) == 32 else "aes-256-cbc"
    result = subprocess.run(
        ["openssl", "enc", f"-{cipher}", "-K", key, "-iv", iv, "-nopad"],
        input=bytes.fromhex(plaintext),
        capture_output=True,
        check=True,
    )
    return result.stdout.hex()


def aes_vectors() -> list[tuple[str, str, str, str]]:
    for key, ciphertext in NIST_CBC:
        assert openssl_cbc(key, NIST_IV, NIST_PT) == ciphertext, key
    return [(key, NIST_IV, NIST_PT, ct) for key, ct in NIST_CBC]


# --------------------------------------------------------------------- Base32

RFC4648_BASE32 = [
    (b"", ""),
    (b"f", "MY======"),
    (b"fo", "MZXQ===="),
    (b"foo", "MZXW6==="),
    (b"foob", "MZXW6YQ="),
    (b"fooba", "MZXW6YTB"),
    (b"foobar", "MZXW6YTBOI======"),
]


def base32_vectors(rng: random.Random) -> list[tuple[str, str]]:
    cases = list(RFC4648_BASE32)
    for data, encoded in cases:
        assert base64.b32encode(data).decode() == encoded
    for length in (10, 16, 20, 32, 64):
        data = bytes(rng.randrange(256) for _ in range(length))
        cases.append((data, base64.b32encode(data).decode()))
    return [(h(d), e) for d, e in cases]


# ------------------------------------------------------------------------ OTP

RFC4226_SECRET = b"12345678901234567890"
RFC4226_CODES = ["755224", "287082", "359152", "969429", "338314", "254676", "287922", "162583", "399871", "520489"]
RFC6238_SEEDS = {
    "SHA1": b"12345678901234567890",
    "SHA256": b"12345678901234567890123456789012",
    "SHA512": b"1234567890" * 6 + b"1234",
}
RFC6238_TIMES = [59, 1111111109, 1111111111, 1234567890, 2000000000, 20000000000]
RFC6238_CODES = {
    "SHA1": ["94287082", "07081804", "14050471", "89005924", "69279037", "65353130"],
    "SHA256": ["46119246", "68084774", "67062674", "91819424", "90698825", "77737706"],
    "SHA512": ["90693936", "25091201", "99943326", "93441116", "38618901", "47863826"],
}


def truncated(alg: str, secret: bytes, counter: int) -> int:
    mac = hmac.new(secret, struct.pack(">Q", counter), alg.lower()).digest()
    offset = mac[-1] & 0x0F
    return int.from_bytes(mac[offset : offset + 4], "big") & 0x7FFFFFFF


def hotp(alg: str, secret: bytes, counter: int, digits: int) -> str:
    return str(truncated(alg, secret, counter) % 10**digits).zfill(digits)


def steam(secret: bytes, counter: int) -> str:
    value = truncated("SHA1", secret, counter)
    chars = []
    for _ in range(5):
        chars.append(STEAM_ALPHABET[value % len(STEAM_ALPHABET)])
        value //= len(STEAM_ALPHABET)
    return "".join(chars)


def otp_vectors(rng: random.Random) -> tuple[list[tuple[int, str, int, int, int, str]], list[tuple[str, int, str]]]:
    """(alg, secret hex, unix time, period, digits, code) and Steam (secret hex, unix time, code).

    HOTP cases use period 0, meaning "time is the counter".
    """
    assert [hotp("SHA1", RFC4226_SECRET, c, 6) for c in range(10)] == RFC4226_CODES
    for alg, codes in RFC6238_CODES.items():
        assert [hotp(alg, RFC6238_SEEDS[alg], t // 30, 8) for t in RFC6238_TIMES] == codes, alg

    cases = [(ALGS["SHA1"], h(RFC4226_SECRET), c, 0, 6, code) for c, code in enumerate(RFC4226_CODES)]
    for alg, codes in RFC6238_CODES.items():
        cases += [
            (ALGS[alg], h(RFC6238_SEEDS[alg]), t, 30, 8, code) for t, code in zip(RFC6238_TIMES, codes, strict=True)
        ]
    boundary_times = [0, 29, 30, 31, 2**31 - 1, 2**31, 2**32 - 1, 4102444800]  # incl. 2038 and the year 2100
    for alg in list(ALGS):
        for _ in range(4):
            secret = bytes(rng.randrange(256) for _ in range(rng.choice((10, 16, 20, 32, 64))))
            for digits in (6, 7, 8):
                period = rng.choice((15, 30, 60, 90, 37))
                for t in (rng.choice(boundary_times), rng.randrange(0, 2**32)):
                    cases.append((ALGS[alg], h(secret), t, period, digits, hotp(alg, secret, t // period, digits)))
    steam_cases = []
    for _ in range(12):
        secret = bytes(rng.randrange(256) for _ in range(20))
        t = rng.randrange(0, 2**32)
        steam_cases.append((h(secret), t, steam(secret, t // 30)))
    return cases, steam_cases


# --------------------------------------------------------------------- output


def mc_value(value: object) -> str:
    if isinstance(value, str):
        return f'"{value}"'
    if isinstance(value, int):
        return f"{value}l" if value > 2**31 - 1 or value < -(2**31) else str(value)
    raise TypeError(value)


def mc_table(name: str, comment: str, rows: Sequence[Sequence[object]]) -> str:
    body = ",\n".join("        [" + ", ".join(mc_value(v) for v in row) + "]" for row in rows)
    return f"    // {comment}\n    const {name} = [\n{body}\n    ];\n"


HEADER = """\
// Author:  cirakg & claude ai
// Summary: Crypto and OTP test vectors (official RFC/FIPS/NIST values plus seeded cases); generated, do not edit.
//
// Generated by tools/gen_vectors.py. Every published vector was recomputed with
// Python's standard library (or openssl for AES) before being written here.
// Algorithm IDs: 0 = SHA-1, 1 = SHA-256, 2 = SHA-512. All data is fake.

import Toybox.Lang;

(:test)
module TestVectors {

"""


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate test/data/Vectors.mc")
    parser.add_argument("--out", type=Path, default=Path("test/data/Vectors.mc"))
    args = parser.parse_args()
    rng = random.Random(SEED)
    otp, steam_cases = otp_vectors(rng)
    parts = [
        mc_table("HASH", "[alg, message hex, digest hex]", list(hash_vectors(rng))),
        mc_table("HMAC", "[alg, key hex, data hex, mac hex] (RFC 2202, RFC 4231)", list(hmac_vectors())),
        mc_table(
            "PBKDF2", "[password hex, salt hex, iterations, length, key hex] (RFC 7914 + seeded)", pbkdf2_vectors()
        ),
        mc_table("HKDF", "[ikm hex, salt hex, info hex, length, okm hex] (RFC 5869 A.1-A.3)", hkdf_vectors()),
        mc_table(
            "AES_CBC", "[key hex, iv hex, plaintext hex, ciphertext hex] (SP 800-38A F.2.1, F.2.5)", aes_vectors()
        ),
        mc_table("BASE32", "[bytes hex, encoded] (RFC 4648 section 10 + seeded)", base32_vectors(rng)),
        mc_table("OTP", "[alg, secret hex, time or counter, period (0 = HOTP), digits, code]", list(otp)),
        mc_table("STEAM", "[secret hex, unix time, code]", steam_cases),
    ]
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(HEADER + "\n".join(parts) + "}\n", encoding="utf-8")
    print(f"wrote {args.out}: {len(otp)} OTP, {len(steam_cases)} Steam vectors")


if __name__ == "__main__":
    main()