Coverage for src/taipanstack/security/models.py: 100%
57 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 15:01 +0000
1"""Secure base models."""
3import json
4import re
5from collections.abc import Callable, Iterator
6from typing import TYPE_CHECKING, Literal, TypeAlias, cast
8from pydantic import BaseModel, ConfigDict
10if TYPE_CHECKING:
11 from pydantic.main import IncEx
12else:
13 IncEx: TypeAlias = set[int] | set[str] | dict[int, object] | dict[str, object]
15from taipanstack.utils.logging import REDACTED_VALUE, SENSITIVE_KEY_PATTERNS
17__all__ = ["SecureBaseModel"]
19_SENSITIVE_KEY_REGEX = (
20 re.compile("|".join(map(re.escape, SENSITIVE_KEY_PATTERNS)), re.IGNORECASE)
21 if SENSITIVE_KEY_PATTERNS
22 else None
23)
25_MAX_RECURSION_DEPTH = 100
28def _mask_dict(data: dict[str, object], depth: int) -> dict[str, object]:
29 """Mask sensitive keys in a dictionary."""
30 masked: dict[str, object] = {}
31 for k, v in data.items():
32 if (
33 isinstance(k, str)
34 and _SENSITIVE_KEY_REGEX is not None
35 and _SENSITIVE_KEY_REGEX.search(k)
36 ):
37 masked[k] = REDACTED_VALUE
38 else:
39 masked[k] = _mask_data(v, depth)
40 return masked
43def _mask_list(data: list[object], depth: int) -> list[object]:
44 """Mask sensitive keys in a list."""
45 return [_mask_data(item, depth) for item in data]
48def _mask_tuple(data: tuple[object, ...], depth: int) -> tuple[object, ...]:
49 """Mask sensitive keys in a tuple."""
50 return tuple(_mask_data(item, depth) for item in data)
53def _mask_set(data: set[object], depth: int) -> set[object]:
54 """Mask sensitive keys in a set."""
55 return {_mask_data(item, depth) for item in data}
58def _mask_collection(data: object, depth: int) -> object:
59 """Dispatch masking based on collection type."""
60 if isinstance(data, dict):
61 return _mask_dict(cast(dict[str, object], data), depth)
62 if isinstance(data, list):
63 return _mask_list(cast(list[object], data), depth)
64 if isinstance(data, tuple):
65 return _mask_tuple(cast(tuple[object, ...], data), depth)
66 if isinstance(data, set):
67 return _mask_set(cast(set[object], data), depth)
68 return data
71def _mask_data(data: object, _depth: int = 0) -> object:
72 """Recursively mask sensitive keys in data."""
73 if _SENSITIVE_KEY_REGEX is None:
74 return data
76 # Prevent ReDoS or stack overflow on deeply nested payloads
77 if _depth > _MAX_RECURSION_DEPTH:
78 return "<MAX_DEPTH_REACHED>"
80 return _mask_collection(data, _depth + 1)
83class SecureBaseModel(BaseModel):
84 """Secure base model that redacts sensitive fields when dumped."""
86 model_config = ConfigDict(frozen=True)
88 def __str__(self) -> str:
89 """Return a string representation with sensitive fields redacted."""
90 return self.__repr__()
92 def __repr_args__(self) -> Iterator[tuple[str | None, object]]:
93 """Provide arguments for string representation, redacting sensitive fields."""
94 for k, v in super().__repr_args__():
95 if (
96 isinstance(k, str)
97 and _SENSITIVE_KEY_REGEX is not None
98 and _SENSITIVE_KEY_REGEX.search(k)
99 ):
100 yield k, REDACTED_VALUE
101 else:
102 yield k, v
104 def model_dump( # noqa: PLR0913
105 self,
106 *,
107 mode: Literal["json", "python"] | str = "python",
108 include: IncEx | None = None,
109 exclude: IncEx | None = None,
110 context: dict[str, object] | None = None,
111 by_alias: bool | None = None,
112 exclude_unset: bool = False,
113 exclude_defaults: bool = False,
114 exclude_none: bool = False,
115 exclude_computed_fields: bool = False,
116 round_trip: bool = False,
117 warnings: bool | Literal["none", "warn", "error"] = True,
118 fallback: Callable[[object], object] | None = None,
119 serialize_as_any: bool = False,
120 **kwargs: object,
121 ) -> dict[str, object]:
122 """Dump the model to a dictionary, redacting sensitive fields.
124 Returns:
125 The redacting dictionary representation of the model.
127 """
128 data = super().model_dump(
129 mode=mode,
130 include=include,
131 exclude=exclude,
132 context=context,
133 by_alias=by_alias,
134 exclude_unset=exclude_unset,
135 exclude_defaults=exclude_defaults,
136 exclude_none=exclude_none,
137 exclude_computed_fields=exclude_computed_fields,
138 round_trip=round_trip,
139 warnings=warnings,
140 fallback=fallback,
141 serialize_as_any=serialize_as_any,
142 **kwargs, # type: ignore[arg-type]
143 )
144 return cast(dict[str, object], _mask_data(data))
146 def model_dump_json( # noqa: PLR0913
147 self,
148 *,
149 indent: int | None = None,
150 ensure_ascii: bool = False,
151 include: IncEx | None = None,
152 exclude: IncEx | None = None,
153 context: dict[str, object] | None = None,
154 by_alias: bool | None = None,
155 exclude_unset: bool = False,
156 exclude_defaults: bool = False,
157 exclude_none: bool = False,
158 exclude_computed_fields: bool = False,
159 round_trip: bool = False,
160 warnings: bool | Literal["none", "warn", "error"] = True,
161 fallback: Callable[[object], object] | None = None,
162 serialize_as_any: bool = False,
163 **kwargs: object,
164 ) -> str:
165 """Dump the model to a JSON string, redacting sensitive fields.
167 Returns:
168 The redacted JSON string representation of the model.
170 """
171 # Extract indent if any, as model_dump does not accept it
173 # Dump to JSON-compatible dict, mask, then serialize
174 dumped_dict = super().model_dump(
175 mode="json",
176 include=include,
177 exclude=exclude,
178 context=context,
179 by_alias=by_alias,
180 exclude_unset=exclude_unset,
181 exclude_defaults=exclude_defaults,
182 exclude_none=exclude_none,
183 exclude_computed_fields=exclude_computed_fields,
184 round_trip=round_trip,
185 warnings=warnings,
186 fallback=fallback,
187 serialize_as_any=serialize_as_any,
188 **kwargs, # type: ignore[arg-type]
189 )
190 masked_dict = _mask_data(dumped_dict)
191 # We need to respect Pydantic's indent/separators if possible,
192 # but json.dumps is the safest standard way.
193 if indent is not None:
194 return json.dumps(masked_dict, indent=indent, ensure_ascii=ensure_ascii)
195 return json.dumps(masked_dict, ensure_ascii=ensure_ascii)