nomyo/SecureCompletionClient.py

@@ -1,4 +1,4 @@
-import json, base64, urllib.parse, httpx, os, secrets, warnings, logging
+import ctypes, json, base64, urllib.parse, httpx, os, secrets, sys, warnings, logging
 from typing import Dict, Any, Optional
 from cryptography.hazmat.primitives import serialization, hashes
 from cryptography.hazmat.primitives.asymmetric import rsa, padding
@@ -354,14 +354,28 @@ class SecureCompletionClient:
             server_public_key_pem.encode('utf-8'),
             backend=default_backend()
         )
-        encrypted_aes_key = server_public_key.encrypt(
+        # RSA encrypt requires bytes — an immutable copy is unavoidable here.
-            aes_key,
+        # We narrow its lifetime to this block and attempt to zero it via
-            padding.OAEP(
+        # CPython internals immediately after use. This relies on the CPython
-                mgf=padding.MGF1(algorithm=hashes.SHA256()),
+        # bytes object layout (ob_sval starts at getsizeof(b'')-1 from id()),
-                algorithm=hashes.SHA256(),
+        # so it is a best-effort measure on CPython only.
-                label=None
+        _key_bytes = bytes(aes_key)
+        try:
+            encrypted_aes_key = server_public_key.encrypt(
+                _key_bytes,
+                padding.OAEP(
+                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                    algorithm=hashes.SHA256(),
+                    label=None
+                )
             )
-        )
+        finally:
+            try:
+                _data_offset = sys.getsizeof(b'') - 1  # offset to ob_sval in PyBytesObject
+                ctypes.memset(id(_key_bytes) + _data_offset, 0, len(_key_bytes))
+            except Exception:
+                pass
+            del _key_bytes
 
         encrypted_package = {
             "version": "1.0",