Fixing tests

aws_lambda_powertools/utilities/data_masking/base.py

@@ -11,14 +11,14 @@ def __init__(self, provider=None):
         else:
             self.provider = provider
 
-    def encrypt(self, data, *args, fields=None, **kwargs):
-        return self._apply_action(data, fields, action=self.provider.encrypt, *args, **kwargs)
+    def encrypt(self, data, fields=None, **kwargs):
+        return self._apply_action(data, fields, self.provider.encrypt, **kwargs)
 
-    def decrypt(self, data, *args, fields=None, **kwargs):
-        return self._apply_action(data, fields, action=self.provider.decrypt, *args, **kwargs)
+    def decrypt(self, data, fields=None, **kwargs):
+        return self._apply_action(data, fields, self.provider.decrypt, **kwargs)
 
-    def mask(self, data, *args, fields=None, **kwargs):
-        return self._apply_action(data, fields, action=self.provider.mask, *args, **kwargs)
+    def mask(self, data, fields=None, **kwargs):
+        return self._apply_action(data, fields, self.provider.mask, **kwargs)
 
     def _apply_action(self, data, fields, action, *args, **kwargs):
         if fields is not None:

aws_lambda_powertools/utilities/data_masking/providers/aws_encryption_sdk.py

@@ -43,30 +43,13 @@ def __init__(self, keys: List[str], client: Optional[EncryptionSDKClient] = None
             max_age=self.MAX_ENTRY_AGE_SECONDS,
             max_messages_encrypted=self.MAX_MESSAGES,
         )
-        self.encryption_context: Dict[str, str] = {}
 
     def encrypt(self, data: Union[bytes, str], *args, **kwargs) -> str:
-        if "context" in kwargs:
-            self.encryption_context = kwargs["context"]
-        ciphertext, header = self.client.encrypt(
-            # Turn all data into string? bc we’re turning everything into a dict
-            # in order to get the key values even if they pass in a json str of a dict
-            source=str(data),
-            encryption_context=self.encryption_context,
-            materials_manager=self.cache_cmm,
-            *args,
-            **kwargs
-        )
+        ciphertext, _ = self.client.encrypt(source=data, materials_manager=self.cache_cmm, *args, **kwargs)
         ciphertext = base64.b64encode(ciphertext).decode()
-        self.encryption_context = header.encryption_context
         return ciphertext
 
     def decrypt(self, data: str, *args, **kwargs) -> bytes:
         ciphertext_decoded = base64.b64decode(data)
-        ciphertext, decrypted_header = self.client.decrypt(
-            source=ciphertext_decoded, key_provider=self.key_provider, *args, **kwargs
-        )
-
-        if self.encryption_context != decrypted_header.encryption_context:
-            raise ValueError("Encryption context mismatch")
+        ciphertext, _ = self.client.decrypt(source=ciphertext_decoded, key_provider=self.key_provider, *args, **kwargs)
         return ciphertext

tests/unit/data_masking/test_data_masking.py

@@ -1,3 +1,4 @@
+import copy
 import json
 
 import pytest
@@ -36,15 +37,29 @@ def decrypt(self, data: str) -> str:
         return serialize.loads(data)
 
 
+data_maskers = [
+    DataMasking(),
+    DataMasking(provider=ItsDangerousProvider("mykey")),
+    DataMasking(provider=AwsEncryptionSdkProvider(keys=[AWS_SDK_KEY])),
+    DataMasking(provider=MyEncryptionProvider(keys="secret-key")),
+]
+
+
 python_dict = {
     "a": {
         "1": {"None": "hello", "four": "world"},  # None type key doesn't work
         "b": {"3": {"4": "goodbye", "e": "world"}},  # key "4.5" doesn't work
     }
 }
 json_dict = json.dumps(python_dict)
-fields = ["a.1.None", "a.b.3.4"]
+dict_fields = ["a.1.None", "a.b.3.4"]
 masked_with_fields = {"a": {"1": {"None": MASK, "four": "world"}, "b": {"3": {"4": MASK, "e": "world"}}}}
+aws_encrypted_with_fields = {
+    "a": {
+        "1": {"None": bytes("hello", "utf-8"), "four": "world"},
+        "b": {"3": {"4": bytes("goodbye", "utf-8"), "e": "world"}},
+    }
+}
 
 # 10kb JSON blob for latency testing
 json_blob = {
@@ -74,8 +89,14 @@ def decrypt(self, data: str) -> str:
     """,
 }
 json_blob_fields = ["address.street", "job_history.company"]
+aws_encrypted_json_blob = copy.deepcopy(json_blob)
+aws_encrypted_json_blob["address"]["street"] = bytes("123 Main St", "utf-8")
+aws_encrypted_json_blob["job_history"]["company"] = bytes("Acme Inc.", "utf-8")
 
-list_of_data_types = [
+dictionaries = [python_dict, json_dict, json_blob]
+fields_to_mask = [dict_fields, dict_fields, json_blob_fields]
+
+data_types_and_masks = [
     # simple data types
     [42, MASK],
     [4.22, MASK],
@@ -85,83 +106,111 @@ def decrypt(self, data: str) -> str:
     # iterables
     [[1, 2, 3, 4], [MASK, MASK, MASK, MASK]],
     [["hello", 1, 2, 3, "world"], [MASK, MASK, MASK, MASK, MASK]],
-    [(55, 66, 88), (MASK, MASK, MASK)],
     # dictionaries
     [python_dict, MASK],
     [json_dict, MASK],
 ]
-
-list_of_data_maskers = [
-    DataMasking(),
-    DataMasking(provider=ItsDangerousProvider("mykey")),
-    DataMasking(provider=AwsEncryptionSdkProvider(keys=[AWS_SDK_KEY])),
-    DataMasking(provider=MyEncryptionProvider(keys="secret-key")),
-]
+data_types = [item[0] for item in data_types_and_masks]
 
 
-@pytest.mark.parametrize("data_masker", list_of_data_maskers)
-@pytest.mark.parametrize("value, value_masked", list_of_data_types)
+@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):
-    """Method to mask a value"""
     masked_string = data_masker.mask(value)
     assert masked_string == value_masked
 
 
-@pytest.mark.parametrize("data_masker", list_of_data_maskers)
-def test_mask_with_fields(data_masker):  # all failing
-    # mask dict with fields
-    masked_string = data_masker.mask(python_dict, fields)
-    print("masked_string in mask with fields: ", masked_string)
+@pytest.mark.parametrize("data_masker", data_maskers)
+def test_mask_with_fields(data_masker):
+    masked_string = data_masker.mask(python_dict, dict_fields)
     assert masked_string == masked_with_fields
-    masked_string = data_masker.mask(json_dict, fields)
-    print("masked_string in mask with fields json: ", masked_string)
+    masked_string = data_masker.mask(json_dict, dict_fields)
     assert masked_string == masked_with_fields
 
 
-@pytest.mark.parametrize("data_masker", list_of_data_maskers)
-@pytest.mark.parametrize("value", list_of_data_types)
+@pytest.mark.parametrize("data_masker", data_maskers)
+@pytest.mark.parametrize("value", data_types)
 def test_encrypt_decrypt(data_masker, value):
-    """Method to encrypt several different data types fully,
-    and specific values in nested dicts"""
-
-    if data_masker == DataMasking():
+    if data_masker == data_maskers[0]:
         with pytest.raises(NotImplementedError):
             encrypted_data = data_masker.encrypt(value)
-            raise NotImplementedError("Subclasses must implement encrypt()")
 
-    encrypted_data = data_masker.encrypt(value)
-    decrypted_data = data_masker.decrypt(encrypted_data)
-    if decrypted_data == [55, 66, 88]:
-        pytest.skip()
-    assert decrypted_data == value
+    else:
+        if data_masker == data_maskers[2]:
+            # AWS Encryption SDK encrypt method only takes in bytes or strings
+            value = bytes(str(value), "utf-8")
 
+        encrypted_data = data_masker.encrypt(value)
+        decrypted_data = data_masker.decrypt(encrypted_data)
+        assert decrypted_data == value
 
-@pytest.mark.parametrize("data_masker", list_of_data_maskers)
-@pytest.mark.parametrize("value", [python_dict, json_dict])
-def test_encrypt_decrypt_with_fields(data_masker, value):
-    if data_masker == DataMasking():
+
+@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)
-            raise NotImplementedError("Subclasses must implement encrypt()")
 
-    encrypted_data = data_masker.encrypt(value, fields)
-    decrypted_data = data_masker.decrypt(encrypted_data, fields)
-    assert decrypted_data == value
+    else:
+        encrypted_data = data_masker.encrypt(value, fields)
+        decrypted_data = data_masker.decrypt(encrypted_data, fields)
 
-    encrypted_json_blob = data_masker.encrypt(json_blob, json_blob_fields)
-    decrypted_json_blob = data_masker.decrypt(encrypted_json_blob, json_blob_fields)
-    assert decrypted_json_blob == json_blob
+        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
 
-
-def test_encrypt_not_implemented():
-    """Test encrypting with no Provider"""
-    data_masker = DataMasking()
-    with pytest.raises(NotImplementedError):
-        data_masker.encrypt("hello world")
+        else:
+            if value == json_dict:
+                assert decrypted_data == json.loads(value)
+            else:
+                assert decrypted_data == value
 
 
 def test_decrypt_not_implemented():
     """Test decrypting with no Provider"""
     data_masker = DataMasking()
     with pytest.raises(NotImplementedError):
         data_masker.decrypt("hello world")
+
+
+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 authentication"}
+    )
+    decrypted_data = data_masker.decrypt(encrypted_data)
+    assert decrypted_data == bytes(str(python_dict), "utf-8")
+
+
+def test_parsing_unsupported_data_type():
+    data_masker = DataMasking()
+    with pytest.raises(TypeError):
+        data_masker.mask(42, ["this.field"])
+
+
+def test_parsing_nonstring_fields():
+    data_masker = DataMasking()
+    _python_dict = {
+        "3": {
+            "1": {"None": "hello", "four": "world"},
+            "4": {"33": {"5": "goodbye", "e": "world"}},
+        }
+    }
+    masked = data_masker.mask(_python_dict, fields=[3.4])
+    assert masked == {"3": {"1": {"None": "hello", "four": "world"}, "4": MASK}}
+
+
+def test_parsing_nonstring_keys_and_fields():
+    data_masker = DataMasking()
+    _python_dict = {
+        3: {
+            "1": {"None": "hello", "four": "world"},
+            4: {"33": {"5": "goodbye", "e": "world"}},
+        }
+    }
+    masked = data_masker.mask(_python_dict, fields=[3.4])
+    assert masked == {"3": {"1": {"None": "hello", "four": "world"}, "4": MASK}}