Updates from PR Feedback

Author: strawgate

key-value/key-value-aio/pyproject.toml

@@ -49,7 +49,7 @@ wrappers-encryption = ["cryptography>=45.0.0"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
-addopts = ["--inline-snapshot=fix,create"]
+addopts = ["--inline-snapshot=disable"]
 markers = [
     "skip_on_ci: Skip running the test when running on CI",
 ]

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

@@ -15,8 +15,8 @@
 _ENCRYPTION_VERSION_KEY = "__encryption_version__"
 _ENCRYPTION_VERSION = 1
 
-EncryptionFn = Callable[[bytes], bytes]
-DecryptionFn = Callable[[bytes], bytes]
+EncryptionFn = Callable[[bytes, int], bytes]
+DecryptionFn = Callable[[bytes, int], bytes]
 
 
 class EncryptionError(Exception):
@@ -45,18 +45,23 @@ def __init__(
         key_value: AsyncKeyValue,
         encryption_fn: EncryptionFn,
         decryption_fn: DecryptionFn,
+        encryption_version: int,
         raise_on_decryption_error: bool = True,
     ) -> None:
         """Initialize the encryption wrapper.
 
         Args:
             key_value: The store to wrap.
-            encryption_fn: The encryption function to use.
-            decryption_fn: The decryption function to use.
+            encryption_fn: The encryption function to use. A callable that takes bytes and an
+                           encryption version int and returns encrypted bytes.
+            decryption_fn: The decryption function to use. A callable that takes bytes and an
+                           encryption version int and returns decrypted bytes.
+            encryption_version: The encryption version to use.
             raise_on_decryption_error: Whether to raise an exception if decryption fails. Defaults to True.
         """
         self.key_value: AsyncKeyValue = key_value
         self.raise_on_decryption_error: bool = raise_on_decryption_error
+        self.encryption_version: int = encryption_version
 
         self._encryption_fn: EncryptionFn = encryption_fn
         self._decryption_fn: DecryptionFn = decryption_fn
@@ -65,9 +70,6 @@ def __init__(
 
     def _encrypt_value(self, value: dict[str, Any]) -> dict[str, Any]:
         """Encrypt a value into the encrypted format."""
-        # # Don't encrypt if it's already encrypted
-        # if _ENCRYPTED_DATA_KEY in value:
-        #     return value
 
         # Serialize to JSON
         try:
@@ -79,14 +81,14 @@ def _encrypt_value(self, value: dict[str, Any]) -> dict[str, Any]:
         json_bytes: bytes = json_str.encode(encoding="utf-8")
 
         # Encrypt with Fernet
-        encrypted_bytes: bytes = self._encryption_fn(json_bytes)
+        encrypted_bytes: bytes = self._encryption_fn(json_bytes, self.encryption_version)
 
         # Encode to base64 for storage in dict (though Fernet output is already base64)
         base64_str: str = base64.b64encode(encrypted_bytes).decode(encoding="ascii")
 
         return {
             _ENCRYPTED_DATA_KEY: base64_str,
-            _ENCRYPTION_VERSION_KEY: _ENCRYPTION_VERSION,
+            _ENCRYPTION_VERSION_KEY: self.encryption_version,
         }
 
     def _decrypt_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
@@ -105,12 +107,18 @@ def _decrypt_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
             msg = f"Corrupted data: expected str, got {type(base64_str)}"
             raise TypeError(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)
+
         try:
             # Decode from base64
             encrypted_bytes: bytes = base64.b64decode(base64_str)
 
             # Decrypt with Fernet
-            json_bytes: bytes = self._decryption_fn(encrypted_bytes)
+            json_bytes: bytes = self._decryption_fn(encrypted_bytes, encryption_version)
 
             # Parse JSON
             json_str: str = json_bytes.decode(encoding="utf-8")

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

@@ -1,9 +1,12 @@
 from cryptography.fernet import Fernet
+from key_value.shared.errors.wrappers.encryption import EncryptionVersionError
 from typing_extensions import overload
 
 from key_value.aio.protocols.key_value import AsyncKeyValue
 from key_value.aio.wrappers.encryption.base import BaseEncryptionWrapper
 
+ENCRYPTION_VERSION = 1
+
 
 class FernetEncryptionWrapper(BaseEncryptionWrapper):
     @overload
@@ -28,7 +31,7 @@ def __init__(
         key_value: AsyncKeyValue,
         *,
         source_material: str,
-        salt: str | None = None,
+        salt: str,
         raise_on_decryption_error: bool = True,
     ) -> None:
         """Initialize the Fernet encryption wrapper.
@@ -57,31 +60,43 @@ def __init__(
             if source_material is None:
                 msg = "Must provide either fernet or source_material"
                 raise ValueError(msg)
+            if salt is None:
+                msg = "Must provide a salt"
+                raise ValueError(msg)
             fernet = Fernet(key=_generate_encryption_key(source_material=source_material, salt=salt))
 
-        def encrypt_with_fernet(data: bytes) -> bytes:
+        def encrypt_with_fernet(data: bytes, encryption_version: int) -> bytes:
+            if encryption_version > self.encryption_version:
+                msg = f"Encryption failed: encryption version {encryption_version} is not supported"
+                raise EncryptionVersionError(msg)
             return fernet.encrypt(data)
 
-        def decrypt_with_fernet(data: bytes) -> bytes:
+        def decrypt_with_fernet(data: bytes, encryption_version: int) -> bytes:
+            if encryption_version > self.encryption_version:
+                msg = f"Decryption failed: encryption version {encryption_version} is not supported"
+                raise EncryptionVersionError(msg)
             return fernet.decrypt(data)
 
         super().__init__(
             key_value=key_value,
             encryption_fn=encrypt_with_fernet,
             decryption_fn=decrypt_with_fernet,
+            encryption_version=ENCRYPTION_VERSION,
             raise_on_decryption_error=raise_on_decryption_error,
         )
 
 
-def _generate_encryption_key(source_material: str, salt: str | None = None) -> bytes:
+def _generate_encryption_key(source_material: str, salt: str) -> bytes:
+    import base64
+
     from cryptography.hazmat.primitives import hashes
     from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 
-    salt = salt or "py-key-value-salt"
-
-    return HKDF(
+    derived_key = HKDF(
         algorithm=hashes.SHA256(),
         length=32,
         salt=salt.encode(),
         info=b"Fernet",
     ).derive(key_material=source_material.encode())
+
+    return base64.urlsafe_b64encode(derived_key)

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

@@ -1,5 +1,6 @@
 import pytest
 from cryptography.fernet import Fernet
+from dirty_equals import IsStr
 from inline_snapshot import snapshot
 from key_value.shared.errors.wrappers.encryption import DecryptionError
 from typing_extensions import override
@@ -40,6 +41,18 @@ async def test_encryption_encrypts_value(self, store: FernetEncryptionWrapper, m
         result = await store.get(collection="test", key="test")
         assert result == original_value
 
+    async def test_encryption_with_wrong_encryption_version(self, store: FernetEncryptionWrapper):
+        """Test that encryption fails with the wrong encryption version."""
+        store.encryption_version = 2
+        original_value = {"test": "value"}
+        await store.put(collection="test", key="test", value=original_value)
+
+        assert await store.get(collection="test", key="test") is not None
+        store.encryption_version = 1
+
+        with pytest.raises(DecryptionError):
+            await store.get(collection="test", key="test")
+
     async def test_encryption_with_string_key(self, store: FernetEncryptionWrapper, memory_store: MemoryStore):
         """Test that encryption works with a string key."""
         original_value = {"test": "value"}
@@ -51,7 +64,7 @@ async def test_encryption_with_string_key(self, store: FernetEncryptionWrapper,
         raw_result = await memory_store.get(collection="test", key="test")
         assert raw_result == snapshot(
             {
-                "__encrypted_data__": "Z0FBQUFBQm8tc0ZsYWhZUmJqUnN0VGlyeGVoUWZuczlPUllyWWxyVEotTVNMVFMtd1hoalNTQk56eFdzNGVocEg0T0xDeEVkTHpJckc2Z0lGZGpCTWZpS3o3cmVWRmRUTl91RENvSW8zNnI3QTlJVmtrQ1FtNnc9",
+                "__encrypted_data__": IsStr(min_length=32),
                 "__encryption_version__": 1,
             }
         )
@@ -109,7 +122,7 @@ 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_wrong_key_returns_original(self, memory_store: MemoryStore):
+    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())
         fernet2 = Fernet(key=Fernet.generate_key())

key-value/key-value-shared/src/key_value/shared/errors/base.py

@@ -5,6 +5,7 @@ class BaseKeyValueError(Exception):
     """Base exception for all KV Store Adapter errors."""
 
     extra_info: ExtraInfoType | None = None
+    message: str | None = None
 
     def __init__(self, message: str | None = None, extra_info: ExtraInfoType | None = None):
         message_parts: list[str] = []
@@ -19,6 +20,8 @@ def __init__(self, message: str | None = None, extra_info: ExtraInfoType | None
 
             message_parts.append(extra_info_str)
 
-        super().__init__(": ".join(message_parts))
+        self.message = ": ".join(message_parts)
+
+        super().__init__(self.message)
 
         self.extra_info = extra_info

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

@@ -7,3 +7,7 @@ class EncryptionError(KeyValueOperationError):
 
 class DecryptionError(EncryptionError):
     """Exception raised when decryption fails."""
+
+
+class EncryptionVersionError(EncryptionError):
+    """Exception raised when the encryption version is not supported."""

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

@@ -18,8 +18,8 @@
 _ENCRYPTION_VERSION_KEY = "__encryption_version__"
 _ENCRYPTION_VERSION = 1
 
-EncryptionFn = Callable[[bytes], bytes]
-DecryptionFn = Callable[[bytes], bytes]
+EncryptionFn = Callable[[bytes, int], bytes]
+DecryptionFn = Callable[[bytes, int], bytes]
 
 
 class EncryptionError(Exception):
@@ -44,18 +44,27 @@ class BaseEncryptionWrapper(BaseWrapper):
     """
 
     def __init__(
-        self, key_value: KeyValue, encryption_fn: EncryptionFn, decryption_fn: DecryptionFn, raise_on_decryption_error: bool = True
+        self,
+        key_value: KeyValue,
+        encryption_fn: EncryptionFn,
+        decryption_fn: DecryptionFn,
+        encryption_version: int,
+        raise_on_decryption_error: bool = True,
     ) -> None:
         """Initialize the encryption wrapper.
 
         Args:
             key_value: The store to wrap.
-            encryption_fn: The encryption function to use.
-            decryption_fn: The decryption function to use.
+            encryption_fn: The encryption function to use. A callable that takes bytes and an
+                           encryption version int and returns encrypted bytes.
+            decryption_fn: The decryption function to use. A callable that takes bytes and an
+                           encryption version int and returns decrypted bytes.
+            encryption_version: The encryption version to use.
             raise_on_decryption_error: Whether to raise an exception if decryption fails. Defaults to True.
         """
         self.key_value: KeyValue = key_value
         self.raise_on_decryption_error: bool = raise_on_decryption_error
+        self.encryption_version: int = encryption_version
 
         self._encryption_fn: EncryptionFn = encryption_fn
         self._decryption_fn: DecryptionFn = decryption_fn
@@ -64,9 +73,6 @@ def __init__(
 
     def _encrypt_value(self, value: dict[str, Any]) -> dict[str, Any]:
         """Encrypt a value into the encrypted format."""
-        # # Don't encrypt if it's already encrypted
-        # if _ENCRYPTED_DATA_KEY in value:
-        #     return value
 
         # Serialize to JSON
         try:
@@ -78,12 +84,12 @@ def _encrypt_value(self, value: dict[str, Any]) -> dict[str, Any]:
         json_bytes: bytes = json_str.encode(encoding="utf-8")
 
         # Encrypt with Fernet
-        encrypted_bytes: bytes = self._encryption_fn(json_bytes)
+        encrypted_bytes: bytes = self._encryption_fn(json_bytes, self.encryption_version)
 
         # Encode to base64 for storage in dict (though Fernet output is already base64)
         base64_str: str = base64.b64encode(encrypted_bytes).decode(encoding="ascii")
 
-        return {_ENCRYPTED_DATA_KEY: base64_str, _ENCRYPTION_VERSION_KEY: _ENCRYPTION_VERSION}
+        return {_ENCRYPTED_DATA_KEY: base64_str, _ENCRYPTION_VERSION_KEY: self.encryption_version}
 
     def _decrypt_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
         """Decrypt a value from the encrypted format."""
@@ -101,12 +107,18 @@ def _decrypt_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
             msg = f"Corrupted data: expected str, got {type(base64_str)}"
             raise TypeError(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)
+
         try:
             # Decode from base64
             encrypted_bytes: bytes = base64.b64decode(base64_str)
 
             # Decrypt with Fernet
-            json_bytes: bytes = self._decryption_fn(encrypted_bytes)
+            json_bytes: bytes = self._decryption_fn(encrypted_bytes, encryption_version)
 
             # Parse JSON
             json_str: str = json_bytes.decode(encoding="utf-8")

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

@@ -2,11 +2,14 @@
 # from the original file 'fernet.py'
 # DO NOT CHANGE! Change the original file instead.
 from cryptography.fernet import Fernet
+from key_value.shared.errors.wrappers.encryption import EncryptionVersionError
 from typing_extensions import overload
 
 from key_value.sync.code_gen.protocols.key_value import KeyValue
 from key_value.sync.code_gen.wrappers.encryption.base import BaseEncryptionWrapper
 
+ENCRYPTION_VERSION = 1
+
 
 class FernetEncryptionWrapper(BaseEncryptionWrapper):
     @overload
@@ -20,9 +23,7 @@ def __init__(self, key_value: KeyValue, *, fernet: Fernet, raise_on_decryption_e
         """
 
     @overload
-    def __init__(
-        self, key_value: KeyValue, *, source_material: str, salt: str | None = None, raise_on_decryption_error: bool = True
-    ) -> None:
+    def __init__(self, key_value: KeyValue, *, source_material: str, salt: str, raise_on_decryption_error: bool = True) -> None:
         """Initialize the Fernet encryption wrapper.
 
         Args:
@@ -49,26 +50,40 @@ def __init__(
             if source_material is None:
                 msg = "Must provide either fernet or source_material"
                 raise ValueError(msg)
+            if salt is None:
+                msg = "Must provide a salt"
+                raise ValueError(msg)
             fernet = Fernet(key=_generate_encryption_key(source_material=source_material, salt=salt))
 
-        def encrypt_with_fernet(data: bytes) -> bytes:
+        def encrypt_with_fernet(data: bytes, encryption_version: int) -> bytes:
+            if encryption_version > self.encryption_version:
+                msg = f"Encryption failed: encryption version {encryption_version} is not supported"
+                raise EncryptionVersionError(msg)
             return fernet.encrypt(data)
 
-        def decrypt_with_fernet(data: bytes) -> bytes:
+        def decrypt_with_fernet(data: bytes, encryption_version: int) -> bytes:
+            if encryption_version > self.encryption_version:
+                msg = f"Decryption failed: encryption version {encryption_version} is not supported"
+                raise EncryptionVersionError(msg)
             return fernet.decrypt(data)
 
         super().__init__(
             key_value=key_value,
             encryption_fn=encrypt_with_fernet,
             decryption_fn=decrypt_with_fernet,
+            encryption_version=ENCRYPTION_VERSION,
             raise_on_decryption_error=raise_on_decryption_error,
         )
 
 
-def _generate_encryption_key(source_material: str, salt: str | None = None) -> bytes:
+def _generate_encryption_key(source_material: str, salt: str) -> bytes:
+    import base64
+
     from cryptography.hazmat.primitives import hashes
     from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 
-    salt = salt or "py-key-value-salt"
+    derived_key = HKDF(algorithm=hashes.SHA256(), length=32, salt=salt.encode(), info=b"Fernet").derive(
+        key_material=source_material.encode()
+    )
 
-    return HKDF(algorithm=hashes.SHA256(), length=32, salt=salt.encode(), info=b"Fernet").derive(key_material=source_material.encode())
+    return base64.urlsafe_b64encode(derived_key)