#!/usr/bin/env python3
"""Derive the vendored ML-DSA-65 test vectors from the ACVP internal projections.

    python3 -I extract.py <ACVP gen-val/json-files directory> <output directory>

Reads `ML-DSA-keyGen-FIPS204/internalProjection.json`, `ML-DSA-sigGen-FIPS204/
internalProjection.json` and `ML-DSA-sigVer-FIPS204/internalProjection.json`, checks each
against the SHA-256 recorded below for the pinned upstream commit, and writes
`ml-dsa-65-keygen.json`, `ml-dsa-65-siggen.json` and `ml-dsa-65-sigver.json`: the same JSON
structure, holding only the test groups whose `parameterSet` is `ML-DSA-65` and which
pq-signatures uses (README.md says which), each test exactly as upstream wrote it, written
with sorted keys, an indent of one space and a final newline. Standard library only.
"""
import hashlib
import json
import os
import sys

UPSTREAM = {
    "ML-DSA-keyGen-FIPS204": "e67ee6540d40e11506c3c4e3b1f79fc1cefcd49820db99fc61f87cc8ba463baf",
    "ML-DSA-sigGen-FIPS204": "72dcaf5f69853ca267ccd16af9cb40949786aca0fcfbf05d1ebeba132b93af22",
    "ML-DSA-sigVer-FIPS204": "47cdd6314c7f746d02421ffcba89d4dbc7bb875ac49e07a029fdfc26fba55437",
}
OUT = {
    "ML-DSA-keyGen-FIPS204": "ml-dsa-65-keygen.json",
    "ML-DSA-sigGen-FIPS204": "ml-dsa-65-siggen.json",
    "ML-DSA-sigVer-FIPS204": "ml-dsa-65-sigver.json",
}


def wanted(mode, g):
    """The groups kept: every ML-DSA-65 key generation group; of signature generation, the
    deterministic groups of the pure external interface and of the internal interface, with
    and without an external mu; of verification, the same three interfaces. The pre-hash
    groups (HashML-DSA, FIPS 204 §5.4) are left out: WAI does not use HashML-DSA."""
    if g.get("parameterSet") != "ML-DSA-65":
        return False
    if mode == "ML-DSA-keyGen-FIPS204":
        return True
    if g.get("preHash") == "preHash":
        return False
    if mode == "ML-DSA-sigGen-FIPS204":
        return g.get("deterministic") is True
    return True


def main(src, dst):
    for mode, digest in UPSTREAM.items():
        path = os.path.join(src, mode, "internalProjection.json")
        with open(path, "rb") as f:
            raw = f.read()
        got = hashlib.sha256(raw).hexdigest()
        if got != digest:
            sys.exit(f"{path}: SHA-256 {got}, expected {digest}")
        doc = json.loads(raw)
        doc["testGroups"] = [g for g in doc["testGroups"] if wanted(mode, g)]
        with open(os.path.join(dst, OUT[mode]), "w", encoding="utf-8") as f:
            json.dump(doc, f, sort_keys=True, indent=1)
            f.write("\n")


if __name__ == "__main__":
    if len(sys.argv) != 3:
        sys.exit(__doc__)
    main(sys.argv[1], sys.argv[2])