Removed args and ItsDangerous and commented on tests

Author: seshubaws

aws_lambda_powertools/utilities/data_masking/base.py

@@ -11,22 +11,49 @@ def __init__(self, provider=None):
         else:
             self.provider = provider
 
-    def encrypt(self, data, fields=None, **kwargs):
-        return self._apply_action(data, fields, self.provider.encrypt, **kwargs)
+    def encrypt(self, data, fields=None, **provider_options):
+        return self._apply_action(data, fields, self.provider.encrypt, **provider_options)
 
-    def decrypt(self, data, fields=None, **kwargs):
-        return self._apply_action(data, fields, self.provider.decrypt, **kwargs)
+    def decrypt(self, data, fields=None, **provider_options):
+        return self._apply_action(data, fields, self.provider.decrypt, **provider_options)
 
-    def mask(self, data, fields=None, **kwargs):
-        return self._apply_action(data, fields, self.provider.mask, **kwargs)
+    def mask(self, data, fields=None, **provider_options):
+        return self._apply_action(data, fields, self.provider.mask, **provider_options)
 
-    def _apply_action(self, data, fields, action, *args, **kwargs):
+    def _apply_action(self, data, fields, action, **provider_options):
         if fields is not None:
-            return self._apply_action_to_fields(data, fields, action, *args, **kwargs)
+            return self._apply_action_to_fields(data, fields, action, **provider_options)
         else:
-            return action(data, *args, **kwargs)
+            return action(data, **provider_options)
+
+    def _apply_action_to_fields(self, data: Union[dict, str], fields, action, **provider_options) -> str:
+        """
+        Apply the specified action to the specified fields in the input data.
+
+        This method is takes the input data, which can be either a dictionary or a JSON string representation
+        of a dictionary, and applies a mask, an encryption, or a decryption to the specified fields.
+
+        Parameters:
+            data (Union[dict, str]): The input data to process. It can be either a dictionary or a JSON string
+                representation of a dictionary.
+            fields (list): A list of fields to apply the action to. Each field can be specified as a string or
+                a list of strings representing nested keys in the dictionary.
+            action (callable): The action to apply to the fields. It should be a callable that takes the current
+                value of the field as the first argument and any additional arguments that might be required
+                for the action. It performs an operation on the current value using the provided arguments and
+                returns the modified value.
+            **provider_options: Additional keyword arguments to pass to the 'action' function.
+
+        Returns:
+            str: A JSON string representation of the modified dictionary after applying the action to the
+            specified fields.
+
+        Raises:
+            ValueError: If 'fields' parameter is None.
+            TypeError: If the 'data' parameter is not a dictionary or a JSON string representation of a dictionary.
+            KeyError: If specified 'fields' do not exist in input data
+        """
 
-    def _apply_action_to_fields(self, data: Union[dict, str], fields, action, *args, **kwargs) -> str:
         if fields is None:
             raise ValueError("No fields specified.")
 
@@ -53,6 +80,6 @@ def _apply_action_to_fields(self, data: Union[dict, str], fields, action, *args,
             for key in keys[:-1]:
                 curr_dict = curr_dict[key]
             valtochange = curr_dict[(keys[-1])]
-            curr_dict[keys[-1]] = action(valtochange, *args, **kwargs)
+            curr_dict[keys[-1]] = action(valtochange, **provider_options)
 
         return my_dict_parsed

aws_lambda_powertools/utilities/data_masking/providers/aws_encryption_sdk.py

@@ -17,9 +17,9 @@ class SingletonMeta(type):
 
     _instances: Dict["AwsEncryptionSdkProvider", Any] = {}
 
-    def __call__(cls, *args, **kwargs):
+    def __call__(cls, *args, **provider_options):
         if cls not in cls._instances:
-            instance = super().__call__(*args, **kwargs)
+            instance = super().__call__(*args, **provider_options)
             cls._instances[cls] = instance
         return cls._instances[cls]
 
@@ -45,12 +45,14 @@ def __init__(self, keys: List[str], client: Optional[EncryptionSDKClient] = None
             max_messages_encrypted=MAX_MESSAGES,
         )
 
-    def encrypt(self, data: Union[bytes, str], *args, **kwargs) -> str:
-        ciphertext, _ = self.client.encrypt(source=data, materials_manager=self.cache_cmm, *args, **kwargs)
+    def encrypt(self, data: Union[bytes, str], **provider_options) -> str:
+        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, *args, **kwargs) -> bytes:
+    def decrypt(self, data: str, **provider_options) -> bytes:
         ciphertext_decoded = base64.b64decode(data)
-        ciphertext, _ = self.client.decrypt(source=ciphertext_decoded, key_provider=self.key_provider, *args, **kwargs)
+        ciphertext, _ = self.client.decrypt(
+            source=ciphertext_decoded, key_provider=self.key_provider, **provider_options
+        )
         return ciphertext

aws_lambda_powertools/utilities/data_masking/providers/itsdangerous.py

@@ -7,14 +7,6 @@
 from aws_lambda_powertools.shared.constants import DATA_MASKING_STRING
 from aws_lambda_powertools.utilities.data_masking.base import DataMasking
 from aws_lambda_powertools.utilities.data_masking.provider import Provider
-from aws_lambda_powertools.utilities.data_masking.providers.aws_encryption_sdk import (
-    AwsEncryptionSdkProvider,
-)
-from aws_lambda_powertools.utilities.data_masking.providers.itsdangerous import (
-    ItsDangerousProvider,
-)
-
-AWS_SDK_KEY = "arn:aws:kms:us-west-2:683517028648:key/269301eb-81eb-4067-ac72-98e8e49bf2b3"
 
 
 class MyEncryptionProvider(Provider):
@@ -39,8 +31,6 @@ def decrypt(self, data: str) -> str:
 
 data_maskers = [
     DataMasking(),
-    DataMasking(provider=ItsDangerousProvider("mykey")),
-    DataMasking(provider=AwsEncryptionSdkProvider(keys=[AWS_SDK_KEY])),
     DataMasking(provider=MyEncryptionProvider(keys="secret-key")),
 ]
 
@@ -121,101 +111,137 @@ def decrypt(self, data: str) -> str:
 @pytest.mark.parametrize("data_masker", data_maskers)
 @pytest.mark.parametrize("value, value_masked", data_types_and_masks)
 def test_mask_types(data_masker, value, value_masked):
+    # GIVEN any data type
+
+    # WHEN mask is called with no fields argument
     masked_string = data_masker.mask(value)
+
+    # THEN the result is the full input data masked
     assert masked_string == value_masked
 
 
 @pytest.mark.parametrize("data_masker", data_maskers)
 def test_mask_with_fields(data_masker):
+    # GIVEN the data type is a dictionary, or a json representation of a dictionary
+
+    # WHEN mask is called with a list of fields specified
     masked_string = data_masker.mask(python_dict, dict_fields)
+    masked_json_string = data_masker.mask(json_dict, dict_fields)
+
+    # THEN the result is only the specified fields are masked
     assert masked_string == masked_with_fields
-    masked_string = data_masker.mask(json_dict, dict_fields)
-    assert masked_string == masked_with_fields
+    assert masked_json_string == masked_with_fields
 
 
-@pytest.mark.parametrize("data_masker", data_maskers)
 @pytest.mark.parametrize("value", data_types)
-def test_encrypt_decrypt(data_masker, value):
-    if data_masker == data_maskers[0]:
-        with pytest.raises(NotImplementedError):
-            encrypted_data = data_masker.encrypt(value)
+def test_encrypt_decrypt(value):
+    # GIVEN an instantiation of DataMasking with a Provider
+    data_masker = DataMasking(provider=MyEncryptionProvider(keys="secret-key"))
 
-    else:
-        if data_masker == data_maskers[2]:
-            # AWS Encryption SDK encrypt method only takes in bytes or strings
-            value = bytes(str(value), "utf-8")
+    # WHEN encrypting and then decrypting the encrypted data
+    encrypted_data = data_masker.encrypt(value)
+    decrypted_data = data_masker.decrypt(encrypted_data)
 
-        encrypted_data = data_masker.encrypt(value)
-        decrypted_data = data_masker.decrypt(encrypted_data)
-        assert decrypted_data == value
+    # THEN the result is the original input data
+    assert decrypted_data == value
 
 
-@pytest.mark.parametrize("data_masker", data_maskers)
 @pytest.mark.parametrize("value, fields", zip(dictionaries, fields_to_mask))
-def test_encrypt_decrypt_with_fields(data_masker, value, fields):
-    if data_masker == data_maskers[0]:
-        with pytest.raises(NotImplementedError):
-            encrypted_data = data_masker.encrypt(value)
+def test_encrypt_decrypt_with_fields(value, fields):
+    # GIVEN an instantiation of DataMasking with a Provider
+    data_masker = DataMasking(provider=MyEncryptionProvider(keys="secret-key"))
+
+    # WHEN encrypting and then decrypting the encrypted data with a list of fields
+    encrypted_data = data_masker.encrypt(value, fields)
+    decrypted_data = data_masker.decrypt(encrypted_data, fields)
 
+    # THEN the result is the original input data
+    if value == json_dict:
+        assert decrypted_data == json.loads(value)
     else:
-        encrypted_data = data_masker.encrypt(value, fields)
-        decrypted_data = data_masker.decrypt(encrypted_data, fields)
+        assert decrypted_data == value
+
+
+def test_encrypt_not_implemented():
+    # GIVEN DataMasking is not initialized with a Provider
+    data_masker = DataMasking()
 
-        if data_masker == data_maskers[2]:
-            # AWS Encryption SDK decrypt method only returns bytes
-            if value == json_blob:
-                assert decrypted_data == aws_encrypted_json_blob
-            else:
-                assert decrypted_data == aws_encrypted_with_fields
+    # WHEN attempting to call the encrypt method on the data
 
-        else:
-            if value == json_dict:
-                assert decrypted_data == json.loads(value)
-            else:
-                assert decrypted_data == value
+    # THEN the result is a NotImplementedError
+    with pytest.raises(NotImplementedError):
+        data_masker.encrypt("hello world")
 
 
 def test_decrypt_not_implemented():
-    """Test decrypting with no Provider"""
+    # GIVEN DataMasking is not initialized with a Provider
     data_masker = DataMasking()
-    with pytest.raises(NotImplementedError):
-        data_masker.decrypt("hello world")
 
+    # WHEN attempting to call the decrypt method on the data
 
-def test_aws_encryption_sdk_with_context():
-    data_masker = DataMasking(provider=AwsEncryptionSdkProvider(keys=[AWS_SDK_KEY]))
-    encrypted_data = data_masker.encrypt(
-        str(python_dict), encryption_context={"not really": "a secret", "but adds": "some auth"}
-    )
-    decrypted_data = data_masker.decrypt(encrypted_data)
-    assert decrypted_data == bytes(str(python_dict), "utf-8")
+    # THEN the result is a NotImplementedError
+    with pytest.raises(NotImplementedError):
+        data_masker.decrypt("hello world")
 
 
 def test_parsing_unsupported_data_type():
+    # GIVEN an initialization of the DataMasking class
     data_masker = DataMasking()
+
+    # WHEN attempting to pass in a list of fields with input data that is not a dict
+
+    # THEN the result is a TypeError
     with pytest.raises(TypeError):
         data_masker.mask(42, ["this.field"])
 
 
+def test_parsing_nonexistent_fields():
+    # GIVEN an initialization of the DataMasking class
+    data_masker = DataMasking()
+    _python_dict = {
+        "3": {
+            "1": {"None": "hello", "four": "world"},
+            "4": {"33": {"5": "goodbye", "e": "world"}},
+        }
+    }
+
+    # WHEN attempting to pass in fields that do not exist in the input data
+
+    # THEN the result is a KeyError
+    with pytest.raises(KeyError):
+        data_masker.mask(_python_dict, ["3.1.True"])
+
+
 def test_parsing_nonstring_fields():
+    # GIVEN an initialization of the DataMasking class
     data_masker = DataMasking()
     _python_dict = {
         "3": {
             "1": {"None": "hello", "four": "world"},
             "4": {"33": {"5": "goodbye", "e": "world"}},
         }
     }
+
+    # WHEN attempting to pass in a list of fields that are not strings
     masked = data_masker.mask(_python_dict, fields=[3.4])
+
+    # THEN the result is the value of the nested field should be masked as normal
     assert masked == {"3": {"1": {"None": "hello", "four": "world"}, "4": DATA_MASKING_STRING}}
 
 
 def test_parsing_nonstring_keys_and_fields():
+    # GIVEN an initialization of the DataMasking class
     data_masker = DataMasking()
+
+    # WHEN the input data is a dictionary with integer keys
     _python_dict = {
         3: {
             "1": {"None": "hello", "four": "world"},
             4: {"33": {"5": "goodbye", "e": "world"}},
         }
     }
+
     masked = data_masker.mask(_python_dict, fields=[3.4])
+
+    # THEN the result is the value of the nested field should be masked as normal
     assert masked == {"3": {"1": {"None": "hello", "four": "world"}, "4": DATA_MASKING_STRING}}