Added custom exception for enc_context mismatch, used pytest fixtures in unit tests

Author: seshubaws

aws_lambda_powertools/utilities/data_masking/providers/aws_encryption_sdk.py

@@ -8,11 +8,16 @@
     LocalCryptoMaterialsCache,
     StrictAwsKmsMasterKeyProvider,
 )
-
 from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
 from aws_lambda_powertools.shared.user_agent import register_feature_to_botocore_session
 
 
+class ContextMismatchError(Exception):
+    def __init__(self, key):
+        super().__init__(f"Encryption Context does not match expected value for key: {key}")
+        self.key = key
+
+
 class SingletonMeta(type):
     """Metaclass to cache class instances to optimize encryption"""
 
@@ -32,27 +37,70 @@ def __call__(cls, *args, **provider_options):
 
 
 class AwsEncryptionSdkProvider(BaseProvider):
-    cache = LocalCryptoMaterialsCache(CACHE_CAPACITY)
+    """
+    The AwsEncryptionSdkProvider is to be used as a Provider for the Datamasking class.
+    Example:
+        >>> data_masker = DataMasking(provider=AwsEncryptionSdkProvider(keys=[keyARN1, keyARN2,...,]))
+        >>> encrypted_data = data_masker.encrypt("a string")
+        "encrptedBase64String"
+        >>> decrypted_data = data_masker.decrypt(encrypted_data)
+        "a string"
+    """
+
     session = botocore.session.Session()
     register_feature_to_botocore_session(session, "data-masking")
 
-    def __init__(self, keys: List[str], client: Optional[EncryptionSDKClient] = None):
+    def __init__(
+        self,
+        keys: List[str],
+        client: Optional[EncryptionSDKClient] = None,
+        local_cache_capacity: Optional[int] = CACHE_CAPACITY,
+        max_cache_age_seconds: Optional[float] = MAX_ENTRY_AGE_SECONDS,
+        max_messages: Optional[int] = MAX_MESSAGES,
+    ):
         self.client = client or EncryptionSDKClient()
         self.keys = keys
+        self.cache = LocalCryptoMaterialsCache(local_cache_capacity)
         self.key_provider = StrictAwsKmsMasterKeyProvider(key_ids=self.keys, botocore_session=self.session)
         self.cache_cmm = CachingCryptoMaterialsManager(
             master_key_provider=self.key_provider,
             cache=self.cache,
-            max_age=MAX_ENTRY_AGE_SECONDS,
-            max_messages_encrypted=MAX_MESSAGES,
+            max_age=max_cache_age_seconds,
+            max_messages_encrypted=max_messages,
         )
 
     def encrypt(self, data: Union[bytes, str], **provider_options) -> str:
+        """
+        Encrypt data using the AwsEncryptionSdkProvider.
+
+        Parameters:
+            - data (Union[bytes, str]):
+                The data to be encrypted.
+            - provider_options:
+                Additional options for the aws_encryption_sdk.EncryptionSDKClient
+
+        Returns:
+            - ciphertext (str):
+                The encrypted data, as a base64-encoded string.
+        """
         ciphertext, _ = self.client.encrypt(source=data, materials_manager=self.cache_cmm, **provider_options)
         ciphertext = base64.b64encode(ciphertext).decode()
         return ciphertext
 
     def decrypt(self, data: str, **provider_options) -> bytes:
+        """
+        Decrypt data using AwsEncryptionSdkProvider.
+
+        Parameters:
+            - data (Union[bytes, str]):
+                The encrypted data, as a base64-encoded string.
+            - provider_options:
+                Additional options for the aws_encryption_sdk.EncryptionSDKClient
+
+        Returns:
+            - ciphertext (bytes):
+                The decrypted data in bytes
+        """
         ciphertext_decoded = base64.b64decode(data)
 
         expected_context = provider_options.pop("encryption_context", {})
@@ -63,6 +111,6 @@ def decrypt(self, data: str, **provider_options) -> bytes:
 
         for key, value in expected_context.items():
             if decryptor_header.encryption_context.get(key) != value:
-                raise ValueError(f"Encryption Context does not match expected value for key: {key}")
+                raise ContextMismatchError(key)
 
         return ciphertext

tests/e2e/data_masking/test_data_masking.py

@@ -4,7 +4,10 @@
 import pytest
 from tests.e2e.utils import data_fetcher
 from aws_lambda_powertools.utilities.data_masking.base import DataMasking
-from aws_lambda_powertools.utilities.data_masking.providers.aws_encryption_sdk import AwsEncryptionSdkProvider
+from aws_lambda_powertools.utilities.data_masking.providers.aws_encryption_sdk import (
+    AwsEncryptionSdkProvider,
+    ContextMismatchError,
+)
 
 
 @pytest.fixture
@@ -52,26 +55,27 @@ def test_encryption_context(data_masker):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider
 
     value = bytes(str([1, 2, "string", 4.5]), "utf-8")
+    context = {"this": "is_secure"}
 
     # WHEN encrypting and then decrypting the encrypted data with an encryption_context
-    encrypted_data = data_masker.encrypt(value, encryption_context={"this": "is_secure"})
-    decrypted_data = data_masker.decrypt(encrypted_data, encryption_context={"this": "is_secure"})
+    encrypted_data = data_masker.encrypt(value, encryption_context=context)
+    decrypted_data = data_masker.decrypt(encrypted_data, encryption_context=context)
 
     # THEN the result is the original input data
     assert decrypted_data == value
 
 
 @pytest.mark.xdist_group(name="data_masking")
-def test_encryption_diff_context_fail(data_masker):
+def test_encryption_context_mismatch(data_masker):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider
 
     value = bytes(str([1, 2, "string", 4.5]), "utf-8")
 
     # WHEN encrypting with a encryption_context
     encrypted_data = data_masker.encrypt(value, encryption_context={"this": "is_secure"})
 
-    # THEN decrypting with a different encryption_context should raise a ValueError
-    with pytest.raises(ValueError):
+    # THEN decrypting with a different encryption_context should raise a ContextMismatchError
+    with pytest.raises(ContextMismatchError):
         data_masker.decrypt(encrypted_data, encryption_context={"not": "same_context"})
 
 
@@ -84,14 +88,14 @@ def test_encryption_no_context_fail(data_masker):
     # WHEN encrypting with no encryption_context
     encrypted_data = data_masker.encrypt(value)
 
-    # THEN decrypting with an encryption_context should raise a ValueError
-    with pytest.raises(ValueError):
+    # THEN decrypting with an encryption_context should raise a ContextMismatchError
+    with pytest.raises(ContextMismatchError):
         data_masker.decrypt(encrypted_data, encryption_context={"this": "is_secure"})
 
 
 # TODO: metaclass?
 @pytest.mark.xdist_group(name="data_masking")
-def test_encryption_key_fail(data_masker, kms_key2_arn):
+def test_encryption_decryption_key_mismatch(data_masker, kms_key2_arn):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider with a certain key
 
     # WHEN encrypting and then decrypting the encrypted data
@@ -106,7 +110,7 @@ def test_encryption_key_fail(data_masker, kms_key2_arn):
 
 
 @pytest.mark.xdist_group(name="data_masking")
-def test_encrypted_in_logs(data_masker, basic_handler_fn, basic_handler_fn_arn):
+def test_encryption_in_logs(data_masker, basic_handler_fn, basic_handler_fn_arn):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider
 
     # WHEN encrypting a value and logging it