Parallélisation de la génération des vecteurs Python

1 files changed, 14 insertions, 3 deletions

diff --git a/test/python/genkat_aead.py b/test/python/genkat_aead.py
index db3a89c..846155f 100755
--- a/test/python/genkat_aead.py
+++ b/test/python/genkat_aead.py
@@ -15,6 +15,9 @@
 
 """Python port of the genkat_aead.c program."""
 
+from os import path
+from sys import argv
+
 import crypto_aead
 
 
@@ -31,7 +34,7 @@ def print_bstr(output, label, buf):
     print('{l} = {b}'.format(l=label, b=buf.hex().upper()), file=output)
 
 
-def generate_test_vectors():
+def generate_test_vectors(output_dir):
     count = 1
     filename = 'LWC_AEAD_KAT_{key}_{npub}.txt'.format(
         key=crypto_aead.KEYBYTES*8, npub=crypto_aead.NPUBBYTES*8
@@ -40,7 +43,7 @@ def generate_test_vectors():
     npub = bytes(range(crypto_aead.NPUBBYTES))
     key = bytes(range(crypto_aead.KEYBYTES))
 
-    with open(filename, 'w') as output:
+    with open(path.join(output_dir, filename), 'w') as output:
 
         for mlen in range(MAX_MESSAGE_LENGTH+1):
             for adlen in range(MAX_ASSOCIATED_DATA_LENGTH+1):
@@ -68,5 +71,13 @@ def generate_test_vectors():
                 print(file=output)
 
 
+def main(argv):
+    output_dir = path.curdir
+    if len(argv) > 1:
+        output_dir = argv[1]
+
+    generate_test_vectors(output_dir)
+
+
 if __name__ == '__main__':
-    generate_test_vectors()
+    main(argv)