add custom mask functionalities

Author: anafalcao

aws_lambda_powertools/utilities/data_masking/base.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import ast
 import functools
 import logging
 import warnings
@@ -94,15 +95,52 @@ def erase(self, data: tuple, fields: list[str]) -> tuple[str]: ...
     @overload
     def erase(self, data: dict, fields: list[str]) -> dict: ...
 
-    def erase(self, data: Sequence | Mapping, fields: list[str] | None = None) -> str | list[str] | tuple[str] | dict:
-        return self._apply_action(data=data, fields=fields, action=self.provider.erase)
+    @overload
+    def erase(
+        self,
+        data: dict,
+        fields: list[str],
+        custom_mask: bool | None = None,
+        mask_pattern: str | None = None,
+        regex_pattern: str | None = None,
+        mask_format: str | None = None,
+    ) -> dict: ...
+
+    def erase(
+        self,
+        data: Sequence | Mapping,
+        fields: list[str] | None = None,
+        custom_mask: bool | None = None,
+        mask_pattern: str | None = None,
+        regex_pattern: str | None = None,
+        mask_format: str | None = None,
+        masking_rules: dict | None = None,
+    ) -> str | list[str] | tuple[str] | dict:
+        if not data:
+            return data
+        if masking_rules:
+            return self._apply_masking_rules(data, masking_rules)
+        else:
+            return self._apply_action(
+                data=data,
+                fields=fields,
+                action=self.provider.erase,
+                custom_mask=custom_mask,
+                mask_pattern=mask_pattern,
+                regex_pattern=regex_pattern,
+                mask_format=mask_format,
+            )
 
     def _apply_action(
         self,
         data,
         fields: list[str] | None,
         action: Callable,
         provider_options: dict | None = None,
+        custom_mask: bool | None = None,
+        mask_pattern: str | None = None,
+        regex_pattern: str | None = None,
+        mask_format: str | None = None,
         **encryption_context: str,
     ):
         """
@@ -136,18 +174,34 @@ def _apply_action(
                 fields=fields,
                 action=action,
                 provider_options=provider_options,
+                custom_mask=custom_mask,
+                mask_pattern=mask_pattern,
+                regex_pattern=regex_pattern,
+                mask_format=mask_format,
                 **encryption_context,
             )
         else:
             logger.debug(f"Running action {action.__name__} with the entire data")
-            return action(data=data, provider_options=provider_options, **encryption_context)
+            return action(
+                data=data,
+                provider_options=provider_options,
+                custom_mask=custom_mask,
+                mask_pattern=mask_pattern,
+                regex_pattern=regex_pattern,
+                mask_format=mask_format,
+                **encryption_context,
+            )
 
     def _apply_action_to_fields(
         self,
         data: dict | str,
         fields: list,
         action: Callable,
         provider_options: dict | None = None,
+        custom_mask: bool | None = None,
+        mask_pattern: str | None = None,
+        regex_pattern: str | None = None,
+        mask_format: str | None = None,
         **encryption_context: str,
     ) -> dict | str:
         """
@@ -194,6 +248,8 @@ def _apply_action_to_fields(
         new_dict = {'a': {'b': {'c': '*****'}}, 'x': {'y': '*****'}}
         ```
         """
+        if not fields:
+            raise ValueError("Fields parameter cannot be empty")
 
         data_parsed: dict = self._normalize_data_to_parse(fields, data)
 
@@ -204,6 +260,10 @@ def _apply_action_to_fields(
             self._call_action,
             action=action,
             provider_options=provider_options,
+            custom_mask=custom_mask,
+            mask_pattern=mask_pattern,
+            regex_pattern=regex_pattern,
+            mask_format=mask_format,
             **encryption_context,  # type: ignore[arg-type]
         )
 
@@ -225,12 +285,6 @@ def _apply_action_to_fields(
             # For in-place updates, json_parse accepts a callback function
             # that receives 3 args: field_value, fields, field_name
             # We create a partial callback to pre-populate known provider options (action, provider opts, enc ctx)
-            update_callback = functools.partial(
-                self._call_action,
-                action=action,
-                provider_options=provider_options,
-                **encryption_context,  # type: ignore[arg-type]
-            )
 
             json_parse.update(
                 data_parsed,
@@ -239,13 +293,60 @@ def _apply_action_to_fields(
 
         return data_parsed
 
+    def _apply_masking_rules(self, data: dict, masking_rules: dict) -> dict:
+        """
+        Apply masking rules to data, supporting different rules for each field.
+        """
+        result = data.copy()
+
+        for path, rule in masking_rules.items():
+            try:
+                # Handle nested paths (e.g., 'address.street')
+                parts = path.split(".")
+                current = result
+
+                for part in parts[:-1]:
+                    if isinstance(current[part], str) and current[part].startswith("{"):
+                        try:
+                            current[part] = ast.literal_eval(current[part])
+                        except (ValueError, SyntaxError):
+                            continue
+                    current = current[part]
+
+                final_field = parts[-1]
+
+                # Apply masking rule to the target field
+                if final_field in current:
+                    current[final_field] = self.provider.erase(str(current[final_field]), **rule)
+
+            except (KeyError, TypeError, AttributeError):
+                # Log warning if field not found or invalid path
+                warnings.warn(f"Could not apply masking rule for path: {path}", stacklevel=2)
+                continue
+
+        return result
+
+    def _mask_nested_field(self, data: dict, field_path: str, mask_function):
+        keys = field_path.split(".")
+        current = data
+        for key in keys[:-1]:
+            current = current.get(key, {})
+            if not isinstance(current, dict):
+                return  # Caminho inválido
+        if keys[-1] in current:
+            current[keys[-1]] = mask_function(current[keys[-1]])
+
     @staticmethod
     def _call_action(
         field_value: Any,
         fields: dict[str, Any],
         field_name: str,
         action: Callable,
         provider_options: dict[str, Any] | None = None,
+        custom_mask: bool | None = None,
+        mask_pattern: str | None = None,
+        regex_pattern: str | None = None,
+        mask_format: str | None = None,
         **encryption_context,
     ) -> None:
         """
@@ -263,7 +364,15 @@ def _call_action(
         Returns:
         - fields[field_name]: Returns the processed field value
         """
-        fields[field_name] = action(field_value, provider_options=provider_options, **encryption_context)
+        fields[field_name] = action(
+            field_value,
+            provider_options=provider_options,
+            custom_mask=custom_mask,
+            mask_pattern=mask_pattern,
+            regex_pattern=regex_pattern,
+            mask_format=mask_format,
+            **encryption_context,
+        )
         return fields[field_name]
 
     def _normalize_data_to_parse(self, fields: list, data: str | dict) -> dict:

aws_lambda_powertools/utilities/data_masking/provider/base.py

@@ -2,10 +2,14 @@
 
 import functools
 import json
+import re
 from typing import Any, Callable, Iterable
 
 from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
 
+PRESERVE_CHARS = set("-_. ")
+_regex_cache = {}
+
 
 class BaseProvider:
     """
@@ -63,7 +67,16 @@ def decrypt(self, data, provider_options: dict | None = None, **encryption_conte
         """
         raise NotImplementedError("Subclasses must implement decrypt()")
 
-    def erase(self, data, **kwargs) -> Iterable[str]:
+    def erase(
+        self,
+        data,
+        custom_mask: bool | None = None,
+        mask_pattern: str | None = None,
+        regex_pattern: str | None = None,
+        mask_format: str | None = None,
+        masking_rules: dict | None = None,
+        **kwargs,
+    ) -> Iterable[str]:
         """
         This method irreversibly erases data.
 
@@ -72,10 +85,58 @@ def erase(self, data, **kwargs) -> Iterable[str]:
 
         If the data to be erased is of an iterable type like `list`, `tuple`,
         or `set`, this method will return a new object of the same type as the
-        input data but with each element replaced by the string "*****".
+        input data but with each element replaced by the string "*****" or following one of the custom masks.
         """
-        if isinstance(data, (str, dict, bytes)):
+        result = DATA_MASKING_STRING
+
+        if data:
+            if isinstance(data, str):
+                if custom_mask:
+                    if mask_pattern:
+                        result = self._pattern_mask(data, mask_pattern)
+                    elif regex_pattern and mask_format:
+                        result = self._regex_mask(data, regex_pattern, mask_format)
+                    else:
+                        result = self._custom_erase(data, **kwargs)
+            elif isinstance(data, dict):
+                if masking_rules:
+                    result = self._apply_masking_rules(data, masking_rules)
+            elif isinstance(data, (list, tuple, set)):
+                result = type(data)(
+                    self.erase(
+                        item,
+                        custom_mask=custom_mask,
+                        mask_pattern=mask_pattern,
+                        regex_pattern=regex_pattern,
+                        mask_format=mask_format,
+                        masking_rules=masking_rules,
+                        **kwargs,
+                    )
+                    for item in data
+                )
+
+        return result
+
+    def _apply_masking_rules(self, data: dict, masking_rules: dict) -> dict:
+        return {
+            key: self.erase(str(value), **masking_rules[key]) if key in masking_rules else str(value)
+            for key, value in data.items()
+        }
+
+    def _pattern_mask(self, data: str, pattern: str) -> str:
+        return pattern[: len(data)] if len(pattern) >= len(data) else pattern
+
+    def _regex_mask(self, data: str, regex_pattern: str, mask_format: str) -> str:
+        try:
+            if regex_pattern not in _regex_cache:
+                _regex_cache[regex_pattern] = re.compile(regex_pattern)
+            return _regex_cache[regex_pattern].sub(mask_format, data)
+        except re.error:
             return DATA_MASKING_STRING
-        elif isinstance(data, (list, tuple, set)):
-            return type(data)([DATA_MASKING_STRING] * len(data))
-        return DATA_MASKING_STRING
+
+    def _custom_erase(self, data: str, **kwargs) -> str:
+        if not data:
+            return ""
+
+        # Use join with list comprehension instead of building list incrementally
+        return "".join("*" if char not in PRESERVE_CHARS else char for char in data)