From 8da34586bd1f3ce7e3079ec08a31cd8672df7d4b Mon Sep 17 00:00:00 2001
From: Justin Patriquin <justinpatriquin@gmail.com>
Date: Fri, 22 Sep 2023 12:21:25 -0300
Subject: [PATCH] add test for crypto.py

---
 pycape/llms/crypto.py      | 12 ++++++------
 pycape/llms/crypto_test.py | 17 +++++++++++++++++
 2 files changed, 23 insertions(+), 6 deletions(-)
 create mode 100644 pycape/llms/crypto_test.py

diff --git a/pycape/llms/crypto.py b/pycape/llms/crypto.py
index a9e3690..5db09f3 100644
--- a/pycape/llms/crypto.py
+++ b/pycape/llms/crypto.py
@@ -11,16 +11,16 @@
 NONCE_SIZE = 12
 
 
-def aes_decrypt(ctxt: bytes, data_key: bytes) -> bytes:
+def aes_decrypt(ctxt: bytes, key: bytes) -> bytes:
     nonce, ctxt = ctxt[:NONCE_SIZE], ctxt[NONCE_SIZE:]
-    encryptor = aead.AESGCM(data_key)
+    encryptor = aead.AESGCM(key)
     ptxt = encryptor.decrypt(nonce, ctxt, None)
     return ptxt
 
 
-def aes_encrypt(aes_key: bytes, ptxt: bytes):
-    encryptor = aead.AESGCM(aes_key)
-    nonce = os.urandom(NONCE_SIZE)  # AESGCM nonce size is 12
+def aes_encrypt(ptxt: bytes, key: bytes):
+    encryptor = aead.AESGCM(key)
+    nonce = os.urandom(NONCE_SIZE)
     ctxt = encryptor.encrypt(nonce, ptxt, None)
     return nonce + ctxt
 
@@ -29,7 +29,7 @@ def envelope_encrypt(public_key: bytes, data: Dict[str, Any]):
     aes_key = os.urandom(32)
     s = json.dumps(data)
 
-    enc_data = aes_encrypt(aes_key, s.encode())
+    enc_data = aes_encrypt(s.encode(), aes_key)
 
     pub = serialization.load_pem_public_key(public_key)
 
diff --git a/pycape/llms/crypto_test.py b/pycape/llms/crypto_test.py
new file mode 100644
index 0000000..c6733ce
--- /dev/null
+++ b/pycape/llms/crypto_test.py
@@ -0,0 +1,17 @@
+import os
+
+from pycape.llms.crypto import aes_decrypt
+from pycape.llms.crypto import aes_encrypt
+
+
+def test_encrypt_decrypt():
+    expected = b"hi there"
+
+    key = os.urandom(32)
+    ciphertext = aes_encrypt(expected, key)
+
+    assert expected == aes_decrypt(ciphertext, key)
+
+
+def test_envelope_encrypt():
+    pass