PR Feedback

Author: strawgate

key-value/key-value-aio/src/key_value/aio/wrappers/encryption/base.py

@@ -4,16 +4,15 @@
 from typing import Any, SupportsFloat
 
 from key_value.shared.errors.key_value import SerializationError
-from key_value.shared.errors.wrappers.encryption import DecryptionError
+from key_value.shared.errors.wrappers.encryption import CorruptedEncryptionDataError, DecryptionError
 from typing_extensions import override
 
 from key_value.aio.protocols.key_value import AsyncKeyValue
 from key_value.aio.wrappers.base import BaseWrapper
 
-# Special keys used to store encrypted data
 _ENCRYPTED_DATA_KEY = "__encrypted_data__"
 _ENCRYPTION_VERSION_KEY = "__encryption_version__"
-_ENCRYPTION_VERSION = 1
+
 
 EncryptionFn = Callable[[bytes], bytes]
 DecryptionFn = Callable[[bytes, int], bytes]
@@ -26,18 +25,9 @@ class EncryptionError(Exception):
 class BaseEncryptionWrapper(BaseWrapper):
     """Wrapper that encrypts values before storing and decrypts on retrieval.
 
-    This wrapper encrypts the JSON-serialized value using Fernet (symmetric encryption)
+    This wrapper encrypts the JSON-serialized value using a custom encryption function
     and stores it as a base64-encoded string within a special key in the dictionary.
     This allows encryption while maintaining the dict[str, Any] interface.
-
-    The encrypted format looks like:
-    {
-        "__encrypted_data__": "base64-encoded-encrypted-data",
-        "__encryption_version__": 1
-    }
-
-    Note: The encryption key must be kept secret and secure. If the key is lost,
-    encrypted data cannot be recovered.
     """
 
     def __init__(
@@ -69,22 +59,31 @@ def __init__(
         super().__init__()
 
     def _encrypt_value(self, value: dict[str, Any]) -> dict[str, Any]:
-        """Encrypt a value into the encrypted format."""
+        """Encrypt a value into the encrypted format.
+
+        The encrypted format looks like:
+        {
+            "__encrypted_data__": "base64-encoded-encrypted-data",
+            "__encryption_version__": 1
+        }
+        """
 
         # Serialize to JSON
         try:
             json_str: str = json.dumps(value, separators=(",", ":"))
+
+            json_bytes: bytes = json_str.encode(encoding="utf-8")
         except (json.JSONDecodeError, TypeError) as e:
             msg: str = f"Failed to serialize object to JSON: {e}"
             raise SerializationError(msg) from e
 
-        json_bytes: bytes = json_str.encode(encoding="utf-8")
-
-        # Encrypt with Fernet
-        encrypted_bytes: bytes = self._encryption_fn(json_bytes)
+        try:
+            encrypted_bytes: bytes = self._encryption_fn(json_bytes)
 
-        # Encode to base64 for storage in dict (though Fernet output is already base64)
-        base64_str: str = base64.b64encode(encrypted_bytes).decode(encoding="ascii")
+            base64_str: str = base64.b64encode(encrypted_bytes).decode(encoding="ascii")
+        except Exception as e:
+            msg = "Failed to encrypt value"
+            raise EncryptionError(msg) from e
 
         return {
             _ENCRYPTED_DATA_KEY: base64_str,
@@ -96,34 +95,30 @@ def _decrypt_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
         if value is None:
             return None
 
-        if _ENCRYPTED_DATA_KEY not in value:
+        if _ENCRYPTED_DATA_KEY not in value and isinstance(value, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
             return value
 
         base64_str = value[_ENCRYPTED_DATA_KEY]
         if not isinstance(base64_str, str):
-            # Corrupted data, return as-is
             msg = f"Corrupted data: expected str, got {type(base64_str)}"
-            raise TypeError(msg)
+            raise CorruptedEncryptionDataError(msg)
 
         if _ENCRYPTION_VERSION_KEY not in value:
             msg = "Corrupted data: missing encryption version"
-            raise TypeError(msg)
+            raise CorruptedEncryptionDataError(msg)
 
         encryption_version = value[_ENCRYPTION_VERSION_KEY]
         if not isinstance(encryption_version, int):
-            # Corrupted data, return as-is
             msg = f"Corrupted data: expected int, got {type(encryption_version)}"
-            raise TypeError(msg)
+            raise CorruptedEncryptionDataError(msg)
 
         try:
-            # Decode from base64
-            encrypted_bytes: bytes = base64.b64decode(base64_str)
+            encrypted_bytes: bytes = base64.b64decode(base64_str, validate=True)
 
-            # Decrypt with Fernet
             json_bytes: bytes = self._decryption_fn(encrypted_bytes, encryption_version)
 
-            # Parse JSON
             json_str: str = json_bytes.decode(encoding="utf-8")
+
             return json.loads(json_str)  # type: ignore[no-any-return]
         except Exception as e:
             msg = "Failed to decrypt value"

key-value/key-value-aio/src/key_value/aio/wrappers/encryption/fernet.py

@@ -1,4 +1,4 @@
-from cryptography.fernet import Fernet
+from cryptography.fernet import Fernet, MultiFernet
 from key_value.shared.errors.wrappers.encryption import EncryptionVersionError
 from typing_extensions import overload
 
@@ -7,21 +7,26 @@
 
 ENCRYPTION_VERSION = 1
 
+KDF_ITERATIONS = 1_200_000
+
 
 class FernetEncryptionWrapper(BaseEncryptionWrapper):
+    """Wrapper that encrypts values before storing and decrypts on retrieval using Fernet (symmetric encryption)."""
+
     @overload
     def __init__(
         self,
         key_value: AsyncKeyValue,
         *,
-        fernet: Fernet,
+        fernet: Fernet | MultiFernet,
         raise_on_decryption_error: bool = True,
     ) -> None:
         """Initialize the Fernet encryption wrapper.
 
         Args:
             key_value: The key-value store to wrap.
-            fernet: The Fernet instance to use for encryption and decryption.
+            fernet: The Fernet or MultiFernet instance to use for encryption and decryption MultiFernet is used to support
+                    key rotation by allowing you to provide multiple Fernet instances that are attempted in order.
             raise_on_decryption_error: Whether to raise an exception if decryption fails. Defaults to True.
         """
 
@@ -47,21 +52,21 @@ def __init__(
         self,
         key_value: AsyncKeyValue,
         *,
-        fernet: Fernet | None = None,
+        fernet: Fernet | MultiFernet | None = None,
         source_material: str | None = None,
         salt: str | None = None,
         raise_on_decryption_error: bool = True,
     ) -> None:
         if fernet is not None:  # noqa: SIM102
             if source_material or salt:
-                msg = "Cannot provide both fernet and source_material and salt"
+                msg = "Cannot provide fernet together with source_material or salt"
                 raise ValueError(msg)
 
         if fernet is None:
-            if not source_material:
+            if not source_material or not source_material.strip():
                 msg = "Must provide either fernet or source_material"
                 raise ValueError(msg)
-            if not salt:
+            if not salt or not salt.strip():
                 msg = "Must provide a salt"
                 raise ValueError(msg)
             fernet = Fernet(key=_generate_encryption_key(source_material=source_material, salt=salt))
@@ -85,7 +90,7 @@ def decrypt_with_fernet(data: bytes, encryption_version: int) -> bytes:
 
 
 def _generate_encryption_key(source_material: str, salt: str) -> bytes:
-    """Generate a Fernet encryption key from a source material and salt using PBKDF2 with 1.2 million iterations."""
+    """Generate a Fernet encryption key from a source material and salt using PBKDF2."""
     import base64
 
     from cryptography.hazmat.primitives import hashes
@@ -95,7 +100,7 @@ def _generate_encryption_key(source_material: str, salt: str) -> bytes:
         algorithm=hashes.SHA256(),
         length=32,
         salt=salt.encode(),
-        iterations=1_200_000,
+        iterations=KDF_ITERATIONS,
     ).derive(key_material=source_material.encode())
 
     return base64.urlsafe_b64encode(pbkdf2)

key-value/key-value-aio/tests/stores/wrappers/test_encryption.py

@@ -1,5 +1,5 @@
 import pytest
-from cryptography.fernet import Fernet
+from cryptography.fernet import Fernet, MultiFernet
 from dirty_equals import IsStr
 from inline_snapshot import snapshot
 from key_value.shared.errors.wrappers.encryption import DecryptionError
@@ -123,6 +123,19 @@ async def test_decryption_ignores_corrupted_data(self, memory_store: MemoryStore
 
         assert await store.get(collection="test", key="test") is None
 
+    async def test_decryption_with_multi_fernet(self, memory_store: MemoryStore):
+        """Test that decryption works with a MultiFernet."""
+        first_fernet = Fernet(key=Fernet.generate_key())
+        first_fernet_store = FernetEncryptionWrapper(key_value=memory_store, fernet=first_fernet)
+        original_value = {"test": "value"}
+        await first_fernet_store.put(collection="test", key="test", value=original_value)
+        assert await first_fernet_store.get(collection="test", key="test") == original_value
+
+        second_fernet = Fernet(key=Fernet.generate_key())
+        multi_fernet = MultiFernet([second_fernet, first_fernet])
+        multi_fernet_store = FernetEncryptionWrapper(key_value=memory_store, fernet=multi_fernet)
+        assert await multi_fernet_store.get(collection="test", key="test") == original_value
+
     async def test_decryption_with_wrong_key_raises_error(self, memory_store: MemoryStore):
         """Test that decryption with the wrong key raises an error."""
         fernet1 = Fernet(key=Fernet.generate_key())

key-value/key-value-shared/src/key_value/shared/errors/wrappers/encryption.py

@@ -11,3 +11,7 @@ class DecryptionError(EncryptionError):
 
 class EncryptionVersionError(EncryptionError):
     """Exception raised when the encryption version is not supported."""
+
+
+class CorruptedEncryptionDataError(EncryptionError):
+    """Exception raised when the encrypted data is corrupted."""

key-value/key-value-shared/src/key_value/shared/utils/managed_entry.py

@@ -85,7 +85,7 @@ def from_json(cls, json_str: str, includes_metadata: bool = True, ttl: SupportsF
 
 def dump_to_json(obj: dict[str, Any]) -> str:
     try:
-        return json.dumps(obj)
+        return json.dumps(obj, separators=(",", ":"))
     except (json.JSONDecodeError, TypeError) as e:
         msg: str = f"Failed to serialize object to JSON: {e}"
         raise SerializationError(msg) from e

key-value/key-value-sync/src/key_value/sync/code_gen/wrappers/encryption/base.py

@@ -7,16 +7,14 @@
 from typing import Any, SupportsFloat
 
 from key_value.shared.errors.key_value import SerializationError
-from key_value.shared.errors.wrappers.encryption import DecryptionError
+from key_value.shared.errors.wrappers.encryption import CorruptedEncryptionDataError, DecryptionError
 from typing_extensions import override
 
 from key_value.sync.code_gen.protocols.key_value import KeyValue
 from key_value.sync.code_gen.wrappers.base import BaseWrapper
 
-# Special keys used to store encrypted data
 _ENCRYPTED_DATA_KEY = "__encrypted_data__"
 _ENCRYPTION_VERSION_KEY = "__encryption_version__"
-_ENCRYPTION_VERSION = 1
 
 EncryptionFn = Callable[[bytes], bytes]
 DecryptionFn = Callable[[bytes, int], bytes]
@@ -29,18 +27,9 @@ class EncryptionError(Exception):
 class BaseEncryptionWrapper(BaseWrapper):
     """Wrapper that encrypts values before storing and decrypts on retrieval.
 
-    This wrapper encrypts the JSON-serialized value using Fernet (symmetric encryption)
+    This wrapper encrypts the JSON-serialized value using a custom encryption function
     and stores it as a base64-encoded string within a special key in the dictionary.
     This allows encryption while maintaining the dict[str, Any] interface.
-
-    The encrypted format looks like:
-    {
-        "__encrypted_data__": "base64-encoded-encrypted-data",
-        "__encryption_version__": 1
-    }
-
-    Note: The encryption key must be kept secret and secure. If the key is lost,
-    encrypted data cannot be recovered.
     """
 
     def __init__(
@@ -72,22 +61,31 @@ def __init__(
         super().__init__()
 
     def _encrypt_value(self, value: dict[str, Any]) -> dict[str, Any]:
-        """Encrypt a value into the encrypted format."""
+        """Encrypt a value into the encrypted format.
+
+        The encrypted format looks like:
+        {
+            "__encrypted_data__": "base64-encoded-encrypted-data",
+            "__encryption_version__": 1
+        }
+        """
 
         # Serialize to JSON
         try:
             json_str: str = json.dumps(value, separators=(",", ":"))
+
+            json_bytes: bytes = json_str.encode(encoding="utf-8")
         except (json.JSONDecodeError, TypeError) as e:
             msg: str = f"Failed to serialize object to JSON: {e}"
             raise SerializationError(msg) from e
 
-        json_bytes: bytes = json_str.encode(encoding="utf-8")
-
-        # Encrypt with Fernet
-        encrypted_bytes: bytes = self._encryption_fn(json_bytes)
+        try:
+            encrypted_bytes: bytes = self._encryption_fn(json_bytes)
 
-        # Encode to base64 for storage in dict (though Fernet output is already base64)
-        base64_str: str = base64.b64encode(encrypted_bytes).decode(encoding="ascii")
+            base64_str: str = base64.b64encode(encrypted_bytes).decode(encoding="ascii")
+        except Exception as e:
+            msg = "Failed to encrypt value"
+            raise EncryptionError(msg) from e
 
         return {_ENCRYPTED_DATA_KEY: base64_str, _ENCRYPTION_VERSION_KEY: self.encryption_version}
 
@@ -96,34 +94,30 @@ def _decrypt_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
         if value is None:
             return None
 
-        if _ENCRYPTED_DATA_KEY not in value:
+        if _ENCRYPTED_DATA_KEY not in value and isinstance(value, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
             return value
 
         base64_str = value[_ENCRYPTED_DATA_KEY]
         if not isinstance(base64_str, str):
-            # Corrupted data, return as-is
             msg = f"Corrupted data: expected str, got {type(base64_str)}"
-            raise TypeError(msg)
+            raise CorruptedEncryptionDataError(msg)
 
         if _ENCRYPTION_VERSION_KEY not in value:
             msg = "Corrupted data: missing encryption version"
-            raise TypeError(msg)
+            raise CorruptedEncryptionDataError(msg)
 
         encryption_version = value[_ENCRYPTION_VERSION_KEY]
         if not isinstance(encryption_version, int):
-            # Corrupted data, return as-is
             msg = f"Corrupted data: expected int, got {type(encryption_version)}"
-            raise TypeError(msg)
+            raise CorruptedEncryptionDataError(msg)
 
         try:
-            # Decode from base64
-            encrypted_bytes: bytes = base64.b64decode(base64_str)
+            encrypted_bytes: bytes = base64.b64decode(base64_str, validate=True)
 
-            # Decrypt with Fernet
             json_bytes: bytes = self._decryption_fn(encrypted_bytes, encryption_version)
 
-            # Parse JSON
             json_str: str = json_bytes.decode(encoding="utf-8")
+
             return json.loads(json_str)  # type: ignore[no-any-return]
         except Exception as e:
             msg = "Failed to decrypt value"