from collections.abc import Mapping from typing import Any, Final from pydantic import BaseModel from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER class SensitiveDataMasker: def __init__( self, sensitive_patterns: set[str] | None = None, non_sensitive_overrides: set[str] | None = None, visible_prefix: int = 4, visible_suffix: int = 4, mask_char: str = "*", mask_short_values: bool = True, ): self.sensitive_patterns = sensitive_patterns or { "password", "secret", "key", "token", "auth", "authorization", "credential", # Plural form: Vertex uses ``vertex_credentials``; segment-exact # matching otherwise misses it because "credential" != "credentials". "credentials", "access", "private", "certificate", "fingerprint", "tenancy", } # If any key segment matches one of these, the key is not considered sensitive # even if it also matches a sensitive pattern. For example, "input_cost_per_token" # contains "token" but "cost" overrides that — it's a pricing field, not a secret. self.non_sensitive_overrides = non_sensitive_overrides or {"cost"} self.visible_prefix = visible_prefix self.visible_suffix = visible_suffix self.mask_char = mask_char self.mask_short_values = mask_short_values def _mask_value(self, value: str) -> str: value_str: Final = str(value) if not value_str: return value if len(value_str) <= (self.visible_prefix + self.visible_suffix): return self.mask_char * len(value_str) if self.mask_short_values else value_str masked_length: Final = len(value_str) - (self.visible_prefix + self.visible_suffix) # Handle the case where visible_suffix is 0 to avoid showing the entire string if self.visible_suffix == 0: return f"{value_str[: self.visible_prefix]}{self.mask_char * masked_length}" else: return ( f"{value_str[: self.visible_prefix]}{self.mask_char * masked_length}{value_str[-self.visible_suffix :]}" ) def is_sensitive_key(self, key: str, excluded_keys: set[str] | None = None) -> bool: # Check if key is in excluded_keys first (exact match) if excluded_keys and key in excluded_keys: return False key_lower: Final = str(key).lower() # Split on underscores/hyphens and check if any segment matches the pattern # This avoids false positives like "max_tokens" matching "token" # but still catches "api_key", "access_token", etc. key_segments: Final = key_lower.replace("-", "_").split("_") # If any segment matches a non-sensitive override, the key is not sensitive. # For example, "input_cost_per_token" contains "token" but also "cost", # so it should not be masked — it's a pricing field, not a secret. if any(override in key_segments for override in self.non_sensitive_overrides): return False result: Final = any(pattern in key_segments for pattern in self.sensitive_patterns) return result def _mask_sequence( self, values: list[Any], depth: int, max_depth: int, excluded_keys: set[str] | None, key_is_sensitive: bool, ) -> list[Any]: masked_items: Final[list[Any]] = [] if depth >= max_depth: return values for item in values: if isinstance(item, Mapping): masked_items.append(self.mask_dict(dict(item), depth + 1, max_depth, excluded_keys)) elif isinstance(item, list): masked_items.append(self._mask_sequence(item, depth + 1, max_depth, excluded_keys, key_is_sensitive)) elif key_is_sensitive and isinstance(item, str): masked_items.append(self._mask_value(item)) else: masked_items.append(item if isinstance(item, (int, float, bool, str, list)) else str(item)) return masked_items def mask_dict( self, data: dict[str, Any], depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, excluded_keys: set[str] | None = None, ) -> dict[str, Any]: if depth >= max_depth: return data masked_data: Final[dict[str, Any]] = {} for k, v in data.items(): try: key_is_sensitive = self.is_sensitive_key(k, excluded_keys) if isinstance(v, Mapping): masked_data[k] = self.mask_dict(dict(v), depth + 1, max_depth, excluded_keys) elif isinstance(v, list): masked_data[k] = self._mask_sequence(v, depth + 1, max_depth, excluded_keys, key_is_sensitive) elif hasattr(v, "__dict__") and not isinstance(v, type): masked_data[k] = self.mask_dict(vars(v), depth + 1, max_depth, excluded_keys) elif key_is_sensitive: str_value = str(v) if v is not None else "" masked_data[k] = self._mask_value(str_value) else: masked_data[k] = v if isinstance(v, (int, float, bool, str, list)) else str(v) except Exception: masked_data[k] = "" return masked_data def mask(self, data: object) -> object: if isinstance(data, Mapping): return self.mask_dict(dict(data)) if isinstance(data, list): return self._mask_sequence( data, 0, DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, None, False, ) return data _default_masker: Final = SensitiveDataMasker() _error_masker: Final = SensitiveDataMasker(visible_prefix=4, visible_suffix=0) def mask_sensitive_structure(data: object) -> object: return _error_masker.mask(data) def mask_credentials_in_payload(data: object) -> object: """Return a copy of ``data`` where string values under sensitive-named keys are masked but every other value (``None``, ``int``, ``float``, ``bool``, ``bytes``, ``datetime``, tuples, sets, typed objects) is preserved by identity, and dicts/lists are rebuilt structurally. Use this for logging payloads that carry response data through to SpendLogs / OTel / Langfuse, where :meth:`SensitiveDataMasker.mask`'s config-dump semantics (``None`` -> ``"None"``, tuples stringified, objects flattened via ``__dict__``) would silently distort the record. Sensitive-key detection is delegated to the shared :class:`SensitiveDataMasker` so pattern updates stay in one place. """ return _walk_payload(data, key_is_sensitive=False, depth=0) def _walk_payload(node: object, key_is_sensitive: bool, depth: int) -> object: if depth >= DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER: return node if isinstance(node, Mapping): return {k: _walk_payload(v, _default_masker.is_sensitive_key(k), depth + 1) for k, v in node.items()} if isinstance(node, list): return [_walk_payload(item, key_is_sensitive, depth + 1) for item in node] if isinstance(node, tuple): return tuple(_walk_payload(item, key_is_sensitive, depth + 1) for item in node) if isinstance(node, BaseModel): return _walk_payload(node.model_dump(), key_is_sensitive, depth) if key_is_sensitive and isinstance(node, str) and node: return _default_masker._mask_value(node) return node def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dict[str, Any]: """Return a new dict with values masked for keys listed in ``sensitive_fields``. Unlike :meth:`SensitiveDataMasker.mask_dict`, this does exact key-name matching (not segment matching), so callers explicitly enumerate which fields to mask. Non-string and None values are passed through unchanged. Values shorter than ``visible_prefix + visible_suffix`` (8 by default) fall outside :meth:`SensitiveDataMasker._mask_value`'s partial-reveal range and are replaced with a fixed-length all-mask string, so a short credential is never returned verbatim. """ masked: Final[dict[str, Any]] = {} mask_char: Final = _default_masker.mask_char min_visible: Final = _default_masker.visible_prefix + _default_masker.visible_suffix for key, value in data.items(): if value is not None and key in sensitive_fields and isinstance(value, str): if len(value) < min_visible: masked[key] = mask_char * len(value) if value else value else: masked[key] = _default_masker._mask_value(value) else: masked[key] = value return masked # Usage example: """ masker = SensitiveDataMasker() data = { "api_key": "sk-1234567890abcdef", "redis_password": "very_secret_pass", "port": 6379, "tags": ["East US 2", "production", "test"] } masked = masker.mask_dict(data) # Result: { # "api_key": "sk-1****cdef", # "redis_password": "very****pass", # "port": 6379, # "tags": ["East US 2", "production", "test"] # } """