1 files changed, 12 insertions, 31 deletions
diff --git a/python/genkat_aead.py b/python/genkat_aead.py
index 8b38d9b..01bed6f 100755
--- a/python/genkat_aead.py
+++ b/python/genkat_aead.py
@@ -1,40 +1,23 @@
 #!/usr/bin/env python3
 
-from lilliput import encrypt, decrypt, LilliputAeMode
-from os import makedirs, path
+import crypto_aead
 
 
 MAX_MESSAGE_LENGTH = 32
 MAX_ADATA_LENGTH = 32
 
-CRYPTO_NPUBBYTES = 120//8
-
-
-MODE_SUFFIXES = {
-    LilliputAeMode.lilliput_1: 'i',
-    LilliputAeMode.lilliput_2: 'ii'
-}
-
 
 def print_bstr(output, label, buf):
     print('{l} = {b}'.format(l=label, b=buf.hex().upper()), file=output)
 
 
-def generate_test_vectors(mode, keylen):
-    print('generating for', mode, keylen)
-
-    directory = 'crypto_aead/lilliputae{mode}{keylen}v1'.format(
-        mode=MODE_SUFFIXES[mode], keylen=keylen
-    )
-
-    makedirs(directory, exist_ok=True)
-
-    output_path = path.join(
-        directory, 'LWC_AEAD_KAT_{keylen}_120.txt'.format(keylen=keylen)
+def generate_test_vectors():
+    output_path = 'LWC_AEAD_KAT_{key}_{npub}.txt'.format(
+        key=crypto_aead.KEYBYTES*8, npub=crypto_aead.NPUBBYTES*8
     )
 
-    nonce = bytes(range(CRYPTO_NPUBBYTES))
-    key = bytes(range(keylen//8))
+    npub = bytes(range(crypto_aead.NPUBBYTES))
+    key = bytes(range(crypto_aead.KEYBYTES))
 
     with open(output_path, 'w') as output:
 
@@ -47,22 +30,20 @@ def generate_test_vectors(mode, keylen):
                 ad = bytes(range(adlen))
 
                 print_bstr(output, 'Key', key)
-                print_bstr(output, 'Nonce', nonce)
+                print_bstr(output, 'Nonce', npub)
                 print_bstr(output, 'PT', msg)
                 print_bstr(output, 'AD', ad)
 
-                ct, tag = encrypt(msg, ad, key, nonce, mode)
+                ct = crypto_aead.encrypt(msg, ad, npub, key)
 
-                print_bstr(output, 'CT', ct+tag)
+                print_bstr(output, 'CT', ct)
 
-                decrypt(ct, tag, ad, key, nonce, mode)
+                crypto_aead.decrypt(ct, ad, npub, key)
 
-                count+=1
+                count += 1
 
                 print(file=output)
 
 
 if __name__ == '__main__':
-    for mode in LilliputAeMode:
-        for keylen in 128, 192, 256:
-            generate_test_vectors(mode, keylen)
+    generate_test_vectors()