Coverage for src/taipanstack/security/models.py: 100%

1"""Secure base models.""" 

2 

3import json 

4import re 

5from collections.abc import Callable, Iterator 

6from typing import TYPE_CHECKING, Literal, TypeAlias, TypeVar, cast 

7 

8from pydantic import BaseModel, ConfigDict 

9 

10if TYPE_CHECKING: 

11 from pydantic.main import IncEx 

12else: 

13 IncEx: TypeAlias = set[int] | set[str] | dict[int, object] | dict[str, object] 

14 

15from taipanstack.utils.logging import REDACTED_VALUE, SENSITIVE_KEY_PATTERNS 

16 

17__all__ = ["SecureBaseModel"] 

18 

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) 

24 

25_MAX_RECURSION_DEPTH = 100 

26 

27 

28def _is_sensitive_key(key: object) -> bool: 

29 """Check if a given key is considered sensitive.""" 

30 return ( 

31 isinstance(key, str) 

32 and _SENSITIVE_KEY_REGEX is not None 

33 and bool(_SENSITIVE_KEY_REGEX.search(key)) 

34 ) 

35 

36 

37VT = TypeVar("VT") 

38 

39 

40def _mask_dict(data: dict[str, VT], depth: int) -> dict[str, object]: 

41 """Mask sensitive keys in a dictionary.""" 

42 masked: dict[str, object] = {} 

43 for k, v in data.items(): 

44 if _is_sensitive_key(k): 

45 masked[k] = REDACTED_VALUE 

46 else: 

47 masked[k] = _mask_data(v, depth) 

48 return masked 

49 

50 

51def _mask_list(data: list[VT], depth: int) -> list[object]: 

52 """Mask sensitive keys in a list.""" 

53 return [_mask_data(item, depth) for item in data] 

54 

55 

56def _mask_tuple(data: tuple[VT, ...], depth: int) -> tuple[object, ...]: 

57 """Mask sensitive keys in a tuple.""" 

58 return tuple(_mask_data(item, depth) for item in data) 

59 

60 

61def _mask_set(data: set[VT], depth: int) -> set[object] | list[object]: 

62 """Mask sensitive keys in a set.""" 

63 masked_items: list[object] = [] 

64 has_unhashable = False 

65 for item in data: 

66 masked = _mask_data(item, depth) 

67 try: 

68 hash(masked) 

69 except TypeError: 

70 has_unhashable = True 

71 masked_items.append(masked) 

72 

73 if has_unhashable: 

74 return masked_items 

75 return set(masked_items) 

76 

77 

78def _mask_sequence(data: object, depth: int) -> object: 

79 """Dispatch masking based on sequence type.""" 

80 if isinstance(data, list): 

81 return _mask_list(cast(list[object], data), depth) 

82 if isinstance(data, tuple): 

83 return _mask_tuple(cast(tuple[object, ...], data), depth) 

84 if isinstance(data, set): 

85 return _mask_set(cast(set[object], data), depth) 

86 return data 

87 

88 

89def _mask_collection(data: object, depth: int) -> object: 

90 """Dispatch masking based on collection type.""" 

91 if isinstance(data, dict): 

92 return _mask_dict(cast(dict[str, object], data), depth) 

93 return _mask_sequence(data, depth) 

94 

95 

96def _mask_data(data: object, _depth: int = 0) -> object: 

97 """Recursively mask sensitive keys in data.""" 

98 if _SENSITIVE_KEY_REGEX is None: 

99 return data 

100 

101 # Prevent ReDoS or stack overflow on deeply nested payloads 

102 if _depth > _MAX_RECURSION_DEPTH: 

103 return "<MAX_DEPTH_REACHED>" 

104 

105 return _mask_collection(data, _depth + 1) 

106 

107 

108class SecureBaseModel(BaseModel): 

109 """Secure base model that redacts sensitive fields when dumped.""" 

110 

111 model_config = ConfigDict(frozen=True) 

112 

113 def __str__(self) -> str: 

114 """Return a string representation with sensitive fields redacted.""" 

115 return self.__repr__() 

116 

117 def __repr_args__(self) -> Iterator[tuple[str | None, object]]: 

118 """Provide arguments for string representation, redacting sensitive fields.""" 

119 for k, v in super().__repr_args__(): 

120 if _is_sensitive_key(k): 

121 yield k, REDACTED_VALUE 

122 else: 

123 yield k, v 

124 

125 def model_dump( # noqa: PLR0913 

126 self, 

127 *, 

128 mode: Literal["json", "python"] | str = "python", 

129 include: IncEx | None = None, 

130 exclude: IncEx | None = None, 

131 context: dict[str, object] | None = None, 

132 by_alias: bool | None = None, 

133 exclude_unset: bool = False, 

134 exclude_defaults: bool = False, 

135 exclude_none: bool = False, 

136 exclude_computed_fields: bool = False, 

137 round_trip: bool = False, 

138 warnings: bool | Literal["none", "warn", "error"] = True, 

139 fallback: Callable[[object], object] | None = None, 

140 serialize_as_any: bool = False, 

141 polymorphic_serialization: bool | None = None, 

142 ) -> dict[str, object]: 

143 """Dump the model to a dictionary, redacting sensitive fields. 

144 

145 Returns: 

146 The redacting dictionary representation of the model. 

147 

148 """ 

149 data = super().model_dump( 

150 mode=mode, 

151 include=include, 

152 exclude=exclude, 

153 context=context, 

154 by_alias=by_alias, 

155 exclude_unset=exclude_unset, 

156 exclude_defaults=exclude_defaults, 

157 exclude_none=exclude_none, 

158 exclude_computed_fields=exclude_computed_fields, 

159 round_trip=round_trip, 

160 warnings=warnings, 

161 fallback=fallback, 

162 serialize_as_any=serialize_as_any, 

163 polymorphic_serialization=polymorphic_serialization, 

164 ) # type: ignore[misc] 

165 return cast(dict[str, object], _mask_data(data)) # type: ignore[misc] 

166 

167 def model_dump_json( # noqa: PLR0913 

168 self, 

169 *, 

170 indent: int | None = None, 

171 ensure_ascii: bool = False, 

172 include: IncEx | None = None, 

173 exclude: IncEx | None = None, 

174 context: dict[str, object] | None = None, 

175 by_alias: bool | None = None, 

176 exclude_unset: bool = False, 

177 exclude_defaults: bool = False, 

178 exclude_none: bool = False, 

179 exclude_computed_fields: bool = False, 

180 round_trip: bool = False, 

181 warnings: bool | Literal["none", "warn", "error"] = True, 

182 fallback: Callable[[object], object] | None = None, 

183 serialize_as_any: bool = False, 

184 polymorphic_serialization: bool | None = None, 

185 ) -> str: 

186 """Dump the model to a JSON string, redacting sensitive fields. 

187 

188 Returns: 

189 The redacted JSON string representation of the model. 

190 

191 """ 

192 # Extract indent if any, as model_dump does not accept it 

193 

194 # Dump to JSON-compatible dict, mask, then serialize 

195 dumped_dict = super().model_dump( 

196 mode="json", 

197 include=include, 

198 exclude=exclude, 

199 context=context, 

200 by_alias=by_alias, 

201 exclude_unset=exclude_unset, 

202 exclude_defaults=exclude_defaults, 

203 exclude_none=exclude_none, 

204 exclude_computed_fields=exclude_computed_fields, 

205 round_trip=round_trip, 

206 warnings=warnings, 

207 fallback=fallback, 

208 serialize_as_any=serialize_as_any, 

209 polymorphic_serialization=polymorphic_serialization, 

210 ) # type: ignore[misc] 

211 masked_dict = _mask_data(dumped_dict) # type: ignore[misc] 

212 # We need to respect Pydantic's indent/separators if possible, 

213 # but json.dumps is the safest standard way. 

214 if indent is not None: 

215 return json.dumps(masked_dict, indent=indent, ensure_ascii=ensure_ascii) 

216 return json.dumps(masked_dict, ensure_ascii=ensure_ascii)