Added serializer for aws enc sdk

Author: seshubaws

aws_lambda_powertools/utilities/data_masking/providers/aws_encryption_sdk.py

@@ -75,6 +75,12 @@ def __init__(
             max_messages_encrypted=max_messages_encrypted,
         )
 
+    def _serialize(self, data: Any):
+        return bytes(str(data), "utf-8")
+
+    def _deserialize(self, data: bytes):
+        return data.decode("utf-8")
+
     def encrypt(self, data: Union[bytes, str], **provider_options) -> str:
         """
         Encrypt data using the AwsEncryptionSdkProvider.
@@ -91,6 +97,7 @@ def encrypt(self, data: Union[bytes, str], **provider_options) -> str:
             ciphertext : str
                 The encrypted data, as a base64-encoded string.
         """
+        data = self._serialize(data)
         ciphertext, _ = self.client.encrypt(source=data, materials_manager=self.cache_cmm, **provider_options)
         ciphertext = base64.b64encode(ciphertext).decode()
         return ciphertext
@@ -125,4 +132,5 @@ def decrypt(self, data: str, **provider_options) -> bytes:
             if decryptor_header.encryption_context.get(key) != value:
                 raise ContextMismatchError(key)
 
+        ciphertext = self._deserialize(ciphertext)
         return ciphertext

tests/e2e/data_masking/handlers/basic_handler.py

@@ -14,7 +14,7 @@ def lambda_handler(event, context):
 
     kms_key = event.get("kms_key")
     data_masker = DataMasking(provider=AwsEncryptionSdkProvider(keys=[kms_key]))
-    value = bytes(str([1, 2, "string", 4.5]), "utf-8")
+    value = [1, 2, "string", 4.5]
     encrypted_data = data_masker.encrypt(value)
     response = {}
     response["encrypted_data"] = encrypted_data

tests/e2e/data_masking/test_data_masking.py

@@ -42,36 +42,36 @@ def test_encryption(data_masker):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider
 
     # AWS Encryption SDK encrypt method only takes in bytes or strings
-    value = bytes(str([1, 2, "string", 4.5]), "utf-8")
+    value = [1, 2, "string", 4.5]
 
     # WHEN encrypting and then decrypting the encrypted data
     encrypted_data = data_masker.encrypt(value)
     decrypted_data = data_masker.decrypt(encrypted_data)
 
     # THEN the result is the original input data
-    assert decrypted_data == value
+    assert decrypted_data == str(value)
 
 
 @pytest.mark.xdist_group(name="data_masking")
 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")
+    value = [1, 2, "string", 4.5]
     context = {"this": "is_secure"}
 
     # WHEN encrypting and then decrypting the encrypted data with an encryption_context
     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
+    assert decrypted_data == str(value)
 
 
 @pytest.mark.xdist_group(name="data_masking")
 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")
+    value = [1, 2, "string", 4.5]
 
     # WHEN encrypting with a encryption_context
     encrypted_data = data_masker.encrypt(value, encryption_context={"this": "is_secure"})
@@ -85,7 +85,7 @@ def test_encryption_context_mismatch(data_masker):
 def test_encryption_no_context_fail(data_masker):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider
 
-    value = bytes(str([1, 2, "string", 4.5]), "utf-8")
+    value = [1, 2, "string", 4.5]
 
     # WHEN encrypting with no encryption_context
     encrypted_data = data_masker.encrypt(value)
@@ -100,7 +100,7 @@ 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
-    value = bytes(str([1, 2, "string", 4.5]), "utf-8")
+    value = [1, 2, "string", 4.5]
     encrypted_data = data_masker.encrypt(value)
 
     # THEN when decrypting with a different key it should fail
@@ -114,12 +114,14 @@ def test_encryption_provider_singleton(data_masker, kms_key1_arn, kms_key2_arn):
     data_masker_2 = DataMasking(provider=AwsEncryptionSdkProvider(keys=[kms_key1_arn]))
     assert data_masker.provider is data_masker_2.provider
 
+    value = [1, 2, "string", 4.5]
+
     # WHEN encrypting and then decrypting the encrypted data
-    encrypted_data = data_masker.encrypt("string")
+    encrypted_data = data_masker.encrypt(value)
     decrypted_data = data_masker_2.decrypt(encrypted_data)
 
     # THEN the result is the original input data
-    assert decrypted_data == bytes("string", "utf-8")
+    assert decrypted_data == str(value)
 
     data_masker_3 = DataMasking(provider=AwsEncryptionSdkProvider(keys=[kms_key2_arn]))
     assert data_masker_2.provider is not data_masker_3.provider
@@ -130,7 +132,7 @@ 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
-    value = bytes(str([1, 2, "string", 4.5]), "utf-8")
+    value = [1, 2, "string", 4.5]
     encrypted_data = data_masker.encrypt(value)
     message = encrypted_data
     custom_key = "order_id"
@@ -146,7 +148,7 @@ def test_encryption_in_logs(data_masker, basic_handler_fn, basic_handler_fn_arn)
     for log in logs.get_log(key=custom_key):
         encrypted_data = log.message
         decrypted_data = data_masker.decrypt(encrypted_data)
-        assert decrypted_data == value
+        assert decrypted_data == str(value)
 
 
 # NOTE: This test is failing currently, need to find a fix for building correct dependencies
@@ -162,4 +164,4 @@ def test_encryption_in_handler(basic_handler_fn_arn, kms_key1_arn):
     decrypted_data = data_masker.decrypt(encrypted_data)
 
     # THEN decrypting the encrypted data from the response should result in the original value
-    assert decrypted_data == bytes(str([1, 2, "string", 4.5]), "utf-8")
+    assert decrypted_data == str([1, 2, "string", 4.5])

tests/functional/data_masking/test_aws_encryption_sdk.py

@@ -39,15 +39,12 @@ def test_mask_with_fields(data_masker):
 def test_encrypt_decrypt(value, data_masker):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider
 
-    # 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)
 
     # THEN the result is the original input data
-    assert decrypted_data == value
+    assert decrypted_data == str(value)
 
 
 @pytest.mark.parametrize("value, fields", zip(dictionaries, fields_to_mask))
@@ -60,7 +57,12 @@ def test_encrypt_decrypt_with_fields(value, fields, data_masker):
 
     # THEN the result is the original input data
     # AWS Encryption SDK decrypt method only returns bytes
+    print("value:", value)
     if value == json_blob:
-        assert decrypted_data == aws_encrypted_json_blob
+        print("json blob!!!!")
+        assert decrypted_data == value
     else:
-        assert decrypted_data == aws_encrypted_with_fields
+        print("json_blob_fields!!!!")
+        assert decrypted_data == str(value)
+        print("decrypted_data:", decrypted_data)
+        print("aws_encrypted_with_fields:", aws_encrypted_with_fields)

tests/unit/data_masking/setup.py

@@ -1,12 +1,8 @@
 import copy
 import json
 
-from aws_lambda_powertools.utilities.data_masking.base import DataMasking
 from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
 
-data_maskers = [DataMasking()]
-
-
 python_dict = {
     "a": {
         "1": {"None": "hello", "four": "world"},  # None type key doesn't work