feat(data_masking): add new sensitive data masking utility

aws_lambda_powertools/shared/constants.py

@@ -6,6 +6,8 @@
 LOGGER_LOG_EVENT_ENV: str = "POWERTOOLS_LOGGER_LOG_EVENT"
 LOGGER_LOG_DEDUPLICATION_ENV: str = "POWERTOOLS_LOG_DEDUPLICATION_DISABLED"
 
+DATA_MASKING_STRING: str = "*****"
+
 MIDDLEWARE_FACTORY_TRACE_ENV: str = "POWERTOOLS_TRACE_MIDDLEWARES"
 
 METRICS_NAMESPACE_ENV: str = "POWERTOOLS_METRICS_NAMESPACE"

aws_lambda_powertools/utilities/data_masking/base.py

@@ -0,0 +1,58 @@
+import json
+from typing import Union
+
+from aws_lambda_powertools.utilities.data_masking.provider import Provider
+
+
+class DataMasking:
+    def __init__(self, provider=None):
+        if provider is None:
+            self.provider = Provider()
+        else:
+            self.provider = provider
+
+    def encrypt(self, data, fields=None, **kwargs):
+        return self._apply_action(data, fields, self.provider.encrypt, **kwargs)
+
+    def decrypt(self, data, fields=None, **kwargs):
+        return self._apply_action(data, fields, self.provider.decrypt, **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:
+            return self._apply_action_to_fields(data, fields, action, *args, **kwargs)
+        else:
+            return action(data, *args, **kwargs)
+
+    def _apply_action_to_fields(self, data: Union[dict, str], fields, action, *args, **kwargs) -> str:
+        if fields is None:
+            raise ValueError("No fields specified.")
+
+        if isinstance(data, str):
+            # Parse JSON string as dictionary
+            my_dict_parsed = json.loads(data)
+        elif isinstance(data, dict):
+            # Turn into json string so everything has quotes around it
+            my_dict_parsed = json.dumps(data)
+            # Turn back into dict so can parse it
+            my_dict_parsed = json.loads(my_dict_parsed)
+        else:
+            raise TypeError(
+                "Unsupported data type. The 'data' parameter must be a dictionary or a JSON string "
+                "representation of a dictionary."
+            )
+
+        for field in fields:
+            if not isinstance(field, str):
+                field = json.dumps(field)
+            keys = field.split(".")
+
+            curr_dict = my_dict_parsed
+            for key in keys[:-1]:
+                curr_dict = curr_dict[key]
+            valtochange = curr_dict[(keys[-1])]
+            curr_dict[keys[-1]] = action(valtochange, *args, **kwargs)
+
+        return my_dict_parsed

aws_lambda_powertools/utilities/data_masking/provider.py

@@ -0,0 +1,26 @@
+from abc import abstractmethod
+from collections.abc import Iterable
+
+from aws_lambda_powertools.shared.constants import DATA_MASKING_STRING
+
+
+class Provider:
+    """
+    When you try to create an instance of a subclass that does not implement the encrypt method,
+    you will get a NotImplementedError with a message that says the method is not implemented:
+    """
+
+    @abstractmethod
+    def encrypt(self, data):
+        raise NotImplementedError("Subclasses must implement encrypt()")
+
+    @abstractmethod
+    def decrypt(self, data):
+        raise NotImplementedError("Subclasses must implement decrypt()")
+
+    def mask(self, data):
+        if isinstance(data, (str, dict, bytes)):
+            return DATA_MASKING_STRING
+        elif isinstance(data, Iterable):
+            return type(data)([DATA_MASKING_STRING] * len(data))
+        return DATA_MASKING_STRING

aws_lambda_powertools/utilities/data_masking/providers/aws_encryption_sdk.py

@@ -0,0 +1,56 @@
+import base64
+from typing import Any, Dict, List, Optional, Union
+
+import botocore
+from aws_encryption_sdk import (
+    CachingCryptoMaterialsManager,
+    EncryptionSDKClient,
+    LocalCryptoMaterialsCache,
+    StrictAwsKmsMasterKeyProvider,
+)
+
+from aws_lambda_powertools.utilities.data_masking.provider import Provider
+
+
+class SingletonMeta(type):
+    """Metaclass to cache class instances to optimize encryption"""
+
+    _instances: Dict["AwsEncryptionSdkProvider", Any] = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            instance = super().__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+CACHE_CAPACITY: int = 100
+MAX_ENTRY_AGE_SECONDS: float = 300.0
+MAX_MESSAGES: int = 200
+# NOTE: You can also set max messages/bytes per data key
+
+
+class AwsEncryptionSdkProvider(Provider, metaclass=SingletonMeta):
+    cache = LocalCryptoMaterialsCache(CACHE_CAPACITY)
+    session = botocore.session.Session()
+
+    def __init__(self, keys: List[str], client: Optional[EncryptionSDKClient] = None) -> None:
+        self.client = client or EncryptionSDKClient()
+        self.keys = keys
+        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,
+        )
+
+    def encrypt(self, data: Union[bytes, str], *args, **kwargs) -> str:
+        ciphertext, _ = self.client.encrypt(source=data, materials_manager=self.cache_cmm, *args, **kwargs)
+        ciphertext = base64.b64encode(ciphertext).decode()
+        return ciphertext
+
+    def decrypt(self, data: str, *args, **kwargs) -> bytes:
+        ciphertext_decoded = base64.b64decode(data)
+        ciphertext, _ = self.client.decrypt(source=ciphertext_decoded, key_provider=self.key_provider, *args, **kwargs)
+        return ciphertext

aws_lambda_powertools/utilities/data_masking/providers/itsdangerous.py

@@ -0,0 +1,53 @@
+from itsdangerous.url_safe import URLSafeSerializer
+
+from aws_lambda_powertools.utilities.data_masking.provider import Provider
+
+
+class ItsDangerousProvider(Provider):
+    def __init__(
+        self,
+        keys,
+        salt=None,
+        serializer=None,
+        serializer_kwargs=None,
+        signer=None,
+        signer_kwargs=None,
+        fallback_signers=None,
+    ):
+        self.keys = keys
+        self.salt = salt
+        self.serializer = serializer
+        self.serializer_kwargs = serializer_kwargs
+        self.signer = signer
+        self.signer_kwargs = signer_kwargs
+        self.fallback_signers = fallback_signers
+
+    def encrypt(self, data):
+        if data is None:
+            return data
+
+        serialized = URLSafeSerializer(
+            self.keys,
+            salt=self.salt,
+            serializer=None,
+            serializer_kwargs=None,
+            signer=None,
+            signer_kwargs=None,
+            fallback_signers=None,
+        )
+        return serialized.dumps(data)
+
+    def decrypt(self, data):
+        if data is None:
+            return data
+
+        serialized = URLSafeSerializer(
+            self.keys,
+            salt=self.salt,
+            serializer=None,
+            serializer_kwargs=None,
+            signer=None,
+            signer_kwargs=None,
+            fallback_signers=None,
+        )
+        return serialized.loads(data)

mypy.ini

@@ -12,6 +12,12 @@ disable_error_code = annotation-unchecked
 [mypy-jmespath]
 ignore_missing_imports=True
 
+[mypy-aws_encryption_sdk]
+ignore_missing_imports=True
+
+[mypy-sentry_sdk]
+ignore_missing_imports=True
+
 [mypy-jmespath.exceptions]
 ignore_missing_imports=True