import hashlib
import sys
import argparse
def err(msg, **kwargs):
print(msg, file=sys.stderr, **kwargs)
def info(msg, **kwargs):
print(msg, **kwargs)
def read_meta_hashes():
hashes = {}
current_name = None
with open("META.yml") as f:
for line in f:
line = line.strip()
if line.startswith("- name:"):
current_name = line.split(":", 1)[1].strip()
elif line.startswith("kat-sha256:") and current_name:
hashes[current_name] = line.split(":", 1)[1].strip()
return hashes
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--scheme",
choices=["44", "65", "87"],
required=True,
help="Parameter set whose KAT output is being checked",
)
args = parser.parse_args()
scheme_name = f"ML-DSA-{args.scheme}"
ref = read_meta_hashes().get(scheme_name)
if ref is None:
err(f"META.yml: no kat-sha256 entry for {scheme_name}")
sys.exit(1)
data = sys.stdin.buffer.read()
if data.startswith(b"SKIPPED"):
info(f"META.yml {scheme_name} kat-sha256: SKIPPED")
sys.exit(0)
computed = hashlib.sha256(data).hexdigest()
if computed == ref:
info(f"META.yml {scheme_name} kat-sha256: OK")
else:
err(f"META.yml {scheme_name} kat-sha256: FAIL ({ref} != {computed})")
sys.exit(1)
if __name__ == "__main__":
main()