From 3e18c92ad2a8c4c7225562d472617ac45b34f8e3 Mon Sep 17 00:00:00 2001
From: Justin Patriquin <justinpatriquin@gmail.com>
Date: Fri, 22 Sep 2023 12:47:40 -0300
Subject: [PATCH] add test for envelope_encryption

---
 pycape/llms/crypto.py      |  2 +-
 pycape/llms/crypto_test.py | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/pycape/llms/crypto.py b/pycape/llms/crypto.py
index 5db09f3..0537c8d 100644
--- a/pycape/llms/crypto.py
+++ b/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)
 
diff --git a/pycape/llms/crypto_test.py b/pycape/llms/crypto_test.py
index c6733ce..4930bea 100644
--- a/pycape/llms/crypto_test.py
+++ b/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)