Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 23 additions & 24 deletions aws_lambda_powertools/utilities/data_masking/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,11 @@
from jsonpath_ng.ext import parse

from aws_lambda_powertools.utilities.data_masking.exceptions import (
DataMaskingError,
DataMaskingFieldNotFoundError,
DataMaskingUnsupportedTypeError,
)
from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
from aws_lambda_powertools.warnings import PowertoolsUserWarning

if TYPE_CHECKING:
from collections.abc import Callable, Mapping, Sequence
Expand Down Expand Up @@ -395,10 +395,10 @@
if not result_parse:
if self.raise_on_missing_field:
# If the data for the field is not found, raise an exception.
raise DataMaskingFieldNotFoundError(f"Field or expression {field_parse} not found in {data_parsed}")
raise DataMaskingFieldNotFoundError(f"Field or expression {field_parse} not found")
else:
# If the data for the field is not found, warning.
warnings.warn(f"Field or expression {field_parse} not found in {data_parsed}", stacklevel=2)
warnings.warn(f"Field or expression {field_parse} not found", stacklevel=2)

# For in-place updates, json_parse accepts a callback function
# that receives 3 args: field_value, fields, field_name
Expand All @@ -411,7 +411,7 @@

return data_parsed

def _apply_masking_rules(self, data: dict, masking_rules: dict) -> dict:

Check failure on line 414 in aws_lambda_powertools/utilities/data_masking/base.py

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Refactor this function to reduce its Cognitive Complexity from 19 to the 15 allowed.

See more on https://sonarcloud.io/project/issues?id=aws-powertools_powertools-lambda-python&issues=AaCA2oArrpSqKOLBLoRA&open=AaCA2oArrpSqKOLBLoRA&pullRequest=8446
"""
Apply masking rules to data, supporting both simple field names and complex path expressions.

Expand All @@ -427,31 +427,30 @@
for path, rule in masking_rules.items():
try:
jsonpath_expr = parse(f"$.{path}")
matches = jsonpath_expr.find(result)
except Exception as exc:
raise DataMaskingError(f"Invalid masking path: {path}") from exc

if not matches:
warnings.warn(f"No matches found for path: {path}", stacklevel=2)
continue
matches = jsonpath_expr.find(result)
if not matches:
if self.raise_on_missing_field:
raise DataMaskingFieldNotFoundError(f"Field or expression {path} not found")

for match in matches:
try:
value = match.value
if value is not None:
masked_value = self.provider.erase(str(value), **rule)
match.full_path.update(result, masked_value)

except Exception as e:
warnings.warn(
f"Error masking value for path {path}: {str(e)}",
category=PowertoolsUserWarning,
stacklevel=2,
)
continue

except Exception as e:
warnings.warn(f"Error processing path {path}: {str(e)}", category=PowertoolsUserWarning, stacklevel=2)
warnings.warn(f"No matches found for path: {path}", stacklevel=2)
continue

for match in matches:
value = match.value
if value is None:
continue

try:
masked_value = self.provider.erase(str(value), **rule)
match.full_path.update(result, masked_value)
except DataMaskingError:
raise
except Exception as exc:
raise DataMaskingError(f"Failed to mask field at path: {path}") from exc

return result

def _mask_nested_field(self, data: dict, field_path: str, mask_function):
Expand Down
18 changes: 12 additions & 6 deletions aws_lambda_powertools/utilities/data_masking/exceptions.py
Original file line number Diff line number Diff line change
@@ -1,34 +1,40 @@
class DataMaskingUnsupportedTypeError(Exception):
class DataMaskingError(Exception):
"""
Base exception for data masking failures.
"""


class DataMaskingUnsupportedTypeError(DataMaskingError):
"""
UnsupportedType Error
"""


class DataMaskingDecryptKeyError(Exception):
class DataMaskingDecryptKeyError(DataMaskingError):
"""
Decrypting with an invalid AWS KMS Key ARN.
"""


class DataMaskingEncryptKeyError(Exception):
class DataMaskingEncryptKeyError(DataMaskingError):
"""
Encrypting with an invalid AWS KMS Key ARN.
"""


class DataMaskingDecryptValueError(Exception):
class DataMaskingDecryptValueError(DataMaskingError):
"""
Decrypting an invalid field.
"""


class DataMaskingContextMismatchError(Exception):
class DataMaskingContextMismatchError(DataMaskingError):
"""
Decrypting with the incorrect encryption context.
"""


class DataMaskingFieldNotFoundError(Exception):
class DataMaskingFieldNotFoundError(DataMaskingError):
"""
Field not found.
"""
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from typing import TYPE_CHECKING, Any

from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
from aws_lambda_powertools.utilities.data_masking.exceptions import DataMaskingError

if TYPE_CHECKING:
from collections.abc import Callable
Expand Down Expand Up @@ -183,8 +184,8 @@ def _regex_mask(self, data: str, regex_pattern: str, mask_format: str) -> str:
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
except re.error as exc:
raise DataMaskingError("Invalid regex pattern") from exc

def _custom_erase(self, data: str) -> str:
if not data:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,10 @@
from aws_lambda_powertools.utilities.data_masking.base import DataMasking
from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
from aws_lambda_powertools.utilities.data_masking.exceptions import (
DataMaskingError,
DataMaskingFieldNotFoundError,
DataMaskingUnsupportedTypeError,
)
from aws_lambda_powertools.warnings import PowertoolsUserWarning


@pytest.fixture
Expand Down Expand Up @@ -315,6 +315,17 @@ def test_no_matches_for_masking_rule(data_masker):
data = {"name": "Ana"}
masking_rules = {"$.missing_field": {"dynamic_mask": True}}

# WHEN applying the masking rule
with pytest.raises(DataMaskingFieldNotFoundError, match=r"Field or expression .* not found"):
data_masker.erase(data=data, masking_rules=masking_rules)


def test_no_matches_for_masking_rule_warning():
# GIVEN a dictionary without the expected field
data_masker = DataMasking(raise_on_missing_field=False)
data = {"name": "Ana"}
masking_rules = {"$.missing_field": {"dynamic_mask": True}}

# WHEN applying the masking rule
with pytest.warns(UserWarning, match=r"No matches found *"):
result = data_masker.erase(data=data, masking_rules=masking_rules)
Expand All @@ -323,7 +334,7 @@ def test_no_matches_for_masking_rule(data_masker):
assert result == data


def test_warning_during_masking_value(data_masker):
def test_error_during_masking_value(data_masker):
# GIVEN data and a masking rule
data = {"value": "test"}

Expand All @@ -335,11 +346,11 @@ def erase(self, value, **kwargs):
data_masker.provider = MockProvider()

# WHEN erase is called
with pytest.warns(expected_warning=PowertoolsUserWarning, match="Error masking value for path value: Mock error"):
masked_data = data_masker.erase(data, masking_rules={"value": {"rule": "value"}})
with pytest.raises(DataMaskingError, match="Failed to mask field at path: value") as exc_info:
data_masker.erase(data, masking_rules={"value": {"rule": "value"}})

# THEN the original data should remain unchanged
assert masked_data["value"] == "test"
# THEN the provider error is preserved
assert isinstance(exc_info.value.__cause__, ValueError)


def test_mask_nested_field_success(data_masker):
Expand Down Expand Up @@ -456,14 +467,12 @@ def test_erase_handles_invalid_regex_pattern(data_masker):
data = "test123"

# WHEN masking with invalid regex
result = data_masker.erase(
data,
regex_pattern="[",
mask_format="X", # Invalid regex pattern that will raise re.error
)

# THEN original data should be returned
assert result == "test123"
with pytest.raises(DataMaskingError, match="Invalid regex pattern"):
data_masker.erase(
data,
regex_pattern="[",
mask_format="X", # Invalid regex pattern that will raise re.error
)


def test_erase_handles_empty_string_with_dynamic_mask(data_masker):
Expand All @@ -485,6 +494,5 @@ def test_erase_dictionary_with_masking_rules_wrong_field(data_masker):
masking_rules = {"user.ssn...": {"custom_mask": "XXX-XX-XXXX"}, "user.address.zip": {"custom_mask": "00000"}}

# WHEN erase is called with wrong masking rules
# We must have a warning
with pytest.warns(expected_warning=PowertoolsUserWarning, match="Error processing path*"):
with pytest.raises(DataMaskingError, match="Invalid masking path"):
data_masker.erase(data, masking_rules=masking_rules)