PR Feedback

Author: strawgate

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

@@ -15,7 +15,7 @@
 _ENCRYPTION_VERSION_KEY = "__encryption_version__"
 _ENCRYPTION_VERSION = 1
 
-EncryptionFn = Callable[[bytes, int], bytes]
+EncryptionFn = Callable[[bytes], bytes]
 DecryptionFn = Callable[[bytes, int], bytes]
 
 
@@ -52,15 +52,15 @@ def __init__(
 
         Args:
             key_value: The store to wrap.
-            encryption_fn: The encryption function to use. A callable that takes bytes and an
-                           encryption version int and returns encrypted bytes.
+            encryption_fn: The encryption function to use. A callable that takes bytes 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
@@ -81,7 +81,7 @@ 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, self.encryption_version)
+        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")
@@ -96,17 +96,19 @@ def _decrypt_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
         if value is None:
             return None
 
-        # Check if it's encrypted
         if _ENCRYPTED_DATA_KEY not in value:
             return value
 
-        # Extract encrypted data
         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)
 
+        if _ENCRYPTION_VERSION_KEY not in value:
+            msg = "Corrupted data: missing encryption version"
+            raise TypeError(msg)
+
         encryption_version = value[_ENCRYPTION_VERSION_KEY]
         if not isinstance(encryption_version, int):
             # Corrupted data, return as-is

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

@@ -52,23 +52,21 @@ def __init__(
         salt: str | None = None,
         raise_on_decryption_error: bool = True,
     ) -> None:
-        if fernet is not None and source_material is not None:
-            msg = "Cannot provide both fernet and source_material"
-            raise ValueError(msg)
+        if fernet is not None:  # noqa: SIM102
+            if source_material or salt:
+                msg = "Cannot provide both fernet and source_material and salt"
+                raise ValueError(msg)
 
         if fernet is None:
-            if source_material is None:
+            if not source_material:
                 msg = "Must provide either fernet or source_material"
                 raise ValueError(msg)
-            if salt is None:
+            if not salt:
                 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, encryption_version: int) -> bytes:
-            if encryption_version > self.encryption_version:
-                msg = f"Encryption failed: encryption version {encryption_version} is not supported"
-                raise EncryptionVersionError(msg)
+        def encrypt_with_fernet(data: bytes) -> bytes:
             return fernet.encrypt(data)
 
         def decrypt_with_fernet(data: bytes, encryption_version: int) -> bytes:

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

@@ -18,7 +18,7 @@
 _ENCRYPTION_VERSION_KEY = "__encryption_version__"
 _ENCRYPTION_VERSION = 1
 
-EncryptionFn = Callable[[bytes, int], bytes]
+EncryptionFn = Callable[[bytes], bytes]
 DecryptionFn = Callable[[bytes, int], bytes]
 
 
@@ -55,15 +55,15 @@ def __init__(
 
         Args:
             key_value: The store to wrap.
-            encryption_fn: The encryption function to use. A callable that takes bytes and an
-                           encryption version int and returns encrypted bytes.
+            encryption_fn: The encryption function to use. A callable that takes bytes 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
@@ -84,7 +84,7 @@ 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, self.encryption_version)
+        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")
@@ -96,17 +96,19 @@ def _decrypt_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
         if value is None:
             return None
 
-        # Check if it's encrypted
         if _ENCRYPTED_DATA_KEY not in value:
             return value
 
-        # Extract encrypted data
         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)
 
+        if _ENCRYPTION_VERSION_KEY not in value:
+            msg = "Corrupted data: missing encryption version"
+            raise TypeError(msg)
+
         encryption_version = value[_ENCRYPTION_VERSION_KEY]
         if not isinstance(encryption_version, int):
             # Corrupted data, return as-is

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

@@ -42,23 +42,21 @@ def __init__(
         salt: str | None = None,
         raise_on_decryption_error: bool = True,
     ) -> None:
-        if fernet is not None and source_material is not None:
-            msg = "Cannot provide both fernet and source_material"
-            raise ValueError(msg)
+        if fernet is not None:  # noqa: SIM102
+            if source_material or salt:
+                msg = "Cannot provide both fernet and source_material and salt"
+                raise ValueError(msg)
 
         if fernet is None:
-            if source_material is None:
+            if not source_material:
                 msg = "Must provide either fernet or source_material"
                 raise ValueError(msg)
-            if salt is None:
+            if not salt:
                 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, encryption_version: int) -> bytes:
-            if encryption_version > self.encryption_version:
-                msg = f"Encryption failed: encryption version {encryption_version} is not supported"
-                raise EncryptionVersionError(msg)
+        def encrypt_with_fernet(data: bytes) -> bytes:
             return fernet.encrypt(data)
 
         def decrypt_with_fernet(data: bytes, encryption_version: int) -> bytes: