add test for envelope_encryption

pycape/llms/crypto.py

@@ -25,7 +25,7 @@ def aes_encrypt(ptxt: bytes, key: bytes):
     return nonce + ctxt
 
 
-def envelope_encrypt(public_key: bytes, data: Dict[str, Any]):
+def envelope_encrypt(public_key: bytes, data: Dict[str, Any]) -> bytes:
     aes_key = os.urandom(32)
     s = json.dumps(data)
 

pycape/llms/crypto_test.py

@@ -1,8 +1,28 @@
+import json
 import os
 
 from pycape.llms.crypto import aes_decrypt
 from pycape.llms.crypto import aes_encrypt
+from pycape.llms.crypto import envelope_encrypt
 
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding, rsa
+
+KEY_PREFIX_LENGTH = 512
+
+def _envelope_decrypt(ciphertext: bytes, priv_key: rsa.RSAPrivateKey):
+    enc_data_key, encrypted_data = ciphertext[:KEY_PREFIX_LENGTH], ciphertext[KEY_PREFIX_LENGTH:]
+
+    data_key = priv_key.decrypt(
+            ciphertext,
+            padding=padding.OAEP(
+                mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None,
+            ),
+        )
+
+    return json.loads(aes_decrypt(encrypted_data, data_key))
 
 def test_encrypt_decrypt():
     expected = b"hi there"
@@ -14,4 +34,14 @@ def test_encrypt_decrypt():
 
 
 def test_envelope_encrypt():
-    pass
+    private_key = rsa.generate_private_key(public_exponent=65537, key_size=4096)
+    pem = private_key.public_key().public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+
+    expected = {"hi": "hello"}
+
+    ciphertext = envelope_encrypt(pem, expected)
+
+    assert expected == _envelope_decrypt(ciphertext, private_key)