add test for crypto.py

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)
 

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