Adapt Kyber KAT generation script for ML-KEM

src/scripts/dev_tools/gen_kyber_kat.py

@@ -1,10 +1,11 @@
 #!/usr/bin/env python3
 
 #
-# Strips the KAT harness in the Kyber reference implementation down
-# to a less space consuming version. This script was used to generate
-# `src/tests/data/pubkey/kyber_kat.vec` test data from the *.rsp files in
-# the reference implementation repository.
+# Strips the KAT harness in the Kyber reference implementation down to a less
+# space consuming version. This script was used to generate:
+#   * `src/tests/data/pubkey/kyber_kat.vec` and
+#   * `src/tests/data/pubkey/ml_kem_ipd.vec`
+# test data from the *.rsp files in the reference implementation repository.
 #
 # See here: https://github.com/pq-crystals/kyber
 #
@@ -29,6 +30,7 @@
 import hashlib
 import binascii
 import os
+import argparse
 
 class KatReader:
     def __init__(self, file):
@@ -97,47 +99,108 @@ def compress_kat(kat, mode_90s):
 
     return kat
 
-def map_mode(mode):
-    out = None
-
-    # Note! See the helper shellscipt in the comment on the top of this file
-    #       to generate KAT *.rsp files that match this naming scheme.
-    if mode == "PQCgenKAT_kem1024-90s":
-        return "Kyber-1024-90s-r3"
-    if mode == "PQCgenKAT_kem512-90s":
-        return "Kyber-512-90s-r3"
-    if mode == "PQCgenKAT_kem768-90s":
-        return "Kyber-768-90s-r3"
-    if mode == "PQCgenKAT_kem1024":
-        return "Kyber-1024-r3"
-    if mode == "PQCgenKAT_kem512":
-        return "Kyber-512-r3"
-    if mode == "PQCgenKAT_kem768":
-        return "Kyber-768-r3"
-
-    raise Exception('Unknown Kyber mode', mode)
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("files", nargs="+", help="Input files")
+    parser.add_argument("--kyber-r3", action="store_true", help="Enable Kyber R3 mode", default=False)
+    parser.add_argument("--ml-kem-ipd", action="store_true", help="Enable ML-KEM initial public draft mode", default=False)
+    parser.add_argument("--ml-kem", action="store_true", help="Enable ML-KEM final mode", default=False)
+    parser.add_argument("--kats-per-mode", type=int, help="Number of KATs to generate per mode", default=25)
+
+    return parser.parse_args()
+
+
+def map_mode(file_name, mode):
+    if mode == "Kyber-r3":
+        # Note! See the helper shellscipt in the comment on the top of this file
+        #       to generate KAT *.rsp files that match this naming scheme.
+        if file_name == "PQCgenKAT_kem1024-90s":
+            return "Kyber-1024-90s-r3"
+        if file_name == "PQCgenKAT_kem512-90s":
+            return "Kyber-512-90s-r3"
+        if file_name == "PQCgenKAT_kem768-90s":
+            return "Kyber-768-90s-r3"
+        if file_name == "PQCgenKAT_kem1024":
+            return "Kyber-1024-r3"
+        if file_name == "PQCgenKAT_kem512":
+            return "Kyber-512-r3"
+        if file_name == "PQCgenKAT_kem768":
+            return "Kyber-768-r3"
+    elif mode == "ML-KEM-ipd":
+        if file_name == "PQCkemKAT_3168":
+            return "ML-KEM-1024-ipd"
+        if file_name == "PQCkemKAT_1632":
+            return "ML-KEM-512-ipd"
+        if file_name == "PQCkemKAT_2400":
+            return "ML-KEM-768-ipd"
+    elif mode == "ML-KEM":
+        if file_name == "PQCkemKAT_3168":
+            return "ML-KEM-1024"
+        if file_name == "PQCkemKAT_1632":
+            return "ML-KEM-512"
+        if file_name == "PQCkemKAT_2400":
+            return "ML-KEM-768"
+    else:
+        raise Exception('Unknown mode', mode)
+
+    raise Exception('Unknown Kyber KAT file name', file_name)
+
+
+def selected_mode(args):
+    modes = []
+
+    if args.kyber_r3:
+        modes.append("Kyber-r3")
+    if args.ml_kem_ipd:
+        modes.append("ML-KEM-ipd")
+    if args.ml_kem:
+        modes.append("ML-KEM")
+
+    if len(modes) > 1:
+        raise Exception("Error: More than one mode selected")
+
+    if len(modes) == 0:
+        raise Exception("Error: No mode selected")
+
+    return modes[0]
+
+
+def output_file(mode):
+    if mode == "Kyber-r3":
+        return "src/tests/data/pubkey/kyber_kat.vec"
+    if mode == "ML-KEM-ipd":
+        return "src/tests/data/pubkey/ml_kem_ipd.vec"
+    if mode == "ML-KEM":
+        return "src/tests/data/pubkey/ml_kem.vec"
+
+    raise Exception("Unknown mode", mode)
+
 
 def main(args = None):
     if args is None:
-        args = sys.argv
+        return 1
+
+    mode = selected_mode(args)
 
-    with open('src/tests/data/pubkey/kyber_kat.vec', 'w') as output:
+    with open(output_file(mode), 'w') as output:
         print("# This file was auto-generated from the reference implemention's KATs", file=output)
         print("# See src/scripts/dev_tools/gen_kyber_kat.py\n", file=output)
 
-        for file in args[1:]:
-            mode = map_mode(os.path.basename(os.path.splitext(file)[0]))
+        for file in args.files:
+            algo_mode = map_mode(os.path.basename(os.path.splitext(file)[0]), mode)
 
             reader = KatReader(open(file))
 
-            print(f"[{mode}]", file=output)
+            print(f"[{algo_mode}]", file=output)
 
-            for kat in list(reader.read_kats())[:25]:
+            for kat in list(reader.read_kats())[:args.kats_per_mode]:
                 kat = compress_kat(kat, '90s' in mode)
 
                 for key in kat.keys():
                     print(key, '=', kat[key], file=output)
                 print("", file=output)
 
 if __name__ == '__main__':
-    sys.exit(main())
+    args = parse_arguments()
+    sys.exit(main(args))