Improvements and fixes

crypto.py

@@ -6,17 +6,34 @@ import os
 from cryptography.hazmat.primitives import padding as sym_padding
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
+from secmem import SecureBuffer, wipe
+
 
 class SymmetricKey:
-    def __init__(self, raw: bytes):
+    def __init__(self, raw: bytes | bytearray):
         if len(raw) != 64:
-            raise ValueError(f"Expected 64 bytes, got {len(raw)}")
-        self.raw = raw
-        self.enc_key = raw[:32]
-        self.mac_key = raw[32:]
+            raise ValueError(f"expected 64 bytes, got {len(raw)}")
+        self._secure = SecureBuffer(raw)
+        if isinstance(raw, bytearray):
+            wipe(raw)
+
+    @property
+    def raw(self) -> bytearray:
+        return self._secure.raw
+
+    @property
+    def enc_key(self) -> bytearray:
+        return self._secure.raw[:32]
+
+    @property
+    def mac_key(self) -> bytearray:
+        return self._secure.raw[32:]
 
     def to_b64(self) -> str:
-        return base64.b64encode(self.raw).decode()
+        return base64.b64encode(bytes(self._secure)).decode()
+
+    def close(self):
+        self._secure.close()
 
     @classmethod
     def from_b64(cls, b64: str) -> "SymmetricKey":
@@ -27,58 +44,61 @@ class SymmetricKey:
         return cls(os.urandom(64))
 
 
+def _decrypt_raw(enc_str: str, key: SymmetricKey) -> bytearray:
+    t, rest = enc_str.split(".", 1)
+    parts = rest.split("|")
+    iv = base64.b64decode(parts[0])
+    ct = base64.b64decode(parts[1])
+    if len(parts) > 2:
+        mac_got = base64.b64decode(parts[2])
+        expected = hmac.new(bytes(key.mac_key), iv + ct, hashlib.sha256).digest()
+        if not hmac.compare_digest(mac_got, expected):
+            raise ValueError("MAC mismatch")
+    dec = Cipher(algorithms.AES(bytes(key.enc_key)), modes.CBC(iv)).decryptor()
+    padded = dec.update(ct) + dec.finalize()
+    unpadder = sym_padding.PKCS7(128).unpadder()
+    return bytearray(unpadder.update(padded) + unpadder.finalize())
+
+
 def enc_string_encrypt(plaintext: str, key: SymmetricKey) -> str:
     iv = os.urandom(16)
-    cipher = Cipher(algorithms.AES(key.enc_key), modes.CBC(iv))
     padder = sym_padding.PKCS7(128).padder()
     padded = padder.update(plaintext.encode()) + padder.finalize()
-    enc = cipher.encryptor()
-    ct = enc.update(padded) + enc.finalize()
-    mac = hmac.new(key.mac_key, iv + ct, hashlib.sha256).digest()
-    return f"2.{base64.b64encode(iv).decode()}|{base64.b64encode(ct).decode()}|{base64.b64encode(mac).decode()}"
+    ct = Cipher(algorithms.AES(bytes(key.enc_key)), modes.CBC(iv)).encryptor()
+    encrypted = ct.update(padded) + ct.finalize()
+    mac = hmac.new(bytes(key.mac_key), iv + encrypted, hashlib.sha256).digest()
+    iv_b64 = base64.b64encode(iv).decode()
+    ct_b64 = base64.b64encode(encrypted).decode()
+    mac_b64 = base64.b64encode(mac).decode()
+    return f"2.{iv_b64}|{ct_b64}|{mac_b64}"
 
 
 def enc_string_decrypt(enc_str: str, key: SymmetricKey) -> str:
-    t, rest = enc_str.split(".", 1)
-    if t != "2":
-        raise ValueError(f"Unsupported type {t}")
-    parts = rest.split("|")
-    iv = base64.b64decode(parts[0])
-    ct = base64.b64decode(parts[1])
-    mac_got = base64.b64decode(parts[2])
-    if not hmac.compare_digest(mac_got, hmac.new(key.mac_key, iv + ct, hashlib.sha256).digest()):
-        raise ValueError("MAC mismatch")
-    dec = Cipher(algorithms.AES(key.enc_key), modes.CBC(iv)).decryptor()
-    padded = dec.update(ct) + dec.finalize()
-    unpadder = sym_padding.PKCS7(128).unpadder()
-    return (unpadder.update(padded) + unpadder.finalize()).decode()
+    raw = _decrypt_raw(enc_str, key)
+    result = raw.decode()
+    wipe(raw)
+    return result
 
 
-def enc_string_decrypt_bytes(enc_str: str, key: SymmetricKey) -> bytes:
-    t, rest = enc_str.split(".", 1)
-    parts = rest.split("|")
-    iv = base64.b64decode(parts[0])
-    ct = base64.b64decode(parts[1])
-    mac_got = base64.b64decode(parts[2]) if len(parts) > 2 else None
-    if mac_got and not hmac.compare_digest(mac_got, hmac.new(key.mac_key, iv + ct, hashlib.sha256).digest()):
-        raise ValueError("MAC mismatch")
-    dec = Cipher(algorithms.AES(key.enc_key), modes.CBC(iv)).decryptor()
-    padded = dec.update(ct) + dec.finalize()
-    unpadder = sym_padding.PKCS7(128).unpadder()
-    return unpadder.update(padded) + unpadder.finalize()
+def enc_string_decrypt_bytes(enc_str: str, key: SymmetricKey) -> bytearray:
+    return _decrypt_raw(enc_str, key)
 
 
 def enc_string_to_dict(enc_str: str) -> dict:
     t, rest = enc_str.split(".", 1)
     parts = rest.split("|")
     d = {"encryptionType": int(t), "encryptedString": enc_str}
-    if len(parts) >= 1: d["iv"] = parts[0]
-    if len(parts) >= 2: d["data"] = parts[1]
-    if len(parts) >= 3: d["mac"] = parts[2]
+    if len(parts) >= 1:
+        d["iv"] = parts[0]
+    if len(parts) >= 2:
+        d["data"] = parts[1]
+    if len(parts) >= 3:
+        d["mac"] = parts[2]
     return d
 
 
 def dict_to_enc_string(d: dict) -> str:
-    if d.get("encryptedString"):
-        return d["encryptedString"]
-    return f"{d.get('encryptionType', 2)}.{d.get('iv', '')}|{d.get('data', '')}|{d.get('mac', '')}"
+    if s := d.get("encryptedString"):
+        return s
+    t = d.get("encryptionType", 2)
+    return f"{t}.{d.get('iv', '')}|{d.get('data', '')}|{d.get('mac', '')}"