From 6881594632606db08d96ed505eed27f2c3978262 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Nov 2025 19:09:48 -0800 Subject: [PATCH] [Fix] Exclude litellm_credential_name from Sensitive Data Masker (Updated) (#16958) * Exclude litellm_credential_name from sensitive masker * Adding missing file --- .../sensitive_data_masker.py | 12 ++--- .../common_utils/openai_endpoint_utils.py | 12 +++-- litellm/proxy/proxy_server.py | 11 +++-- .../test_sensitive_data_masker.py | 46 +++++++++++++++++++ .../test_openai_endpoint_utils.py | 34 ++++++++++++++ 5 files changed, 103 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index c6906b26236..206810943ca 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -42,7 +42,11 @@ class SensitiveDataMasker: 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) -> bool: + def is_sensitive_key(self, key: str, excluded_keys: Optional[Set[str]] = None) -> bool: + # Check if key is in excluded_keys first (exact match) + if excluded_keys and key in excluded_keys: + return False + key_lower = str(key).lower() # Split on underscores and check if any segment matches the pattern # This avoids false positives like "max_tokens" matching "token" @@ -64,7 +68,6 @@ class SensitiveDataMasker: if depth >= max_depth: return data - excluded_keys = excluded_keys or set() masked_data: Dict[str, Any] = {} for k, v in data.items(): try: @@ -72,10 +75,7 @@ class SensitiveDataMasker: masked_data[k] = self.mask_dict(v, depth + 1, max_depth, excluded_keys) elif hasattr(v, "__dict__") and not isinstance(v, type): masked_data[k] = self.mask_dict(vars(v), depth + 1, max_depth, excluded_keys) - elif k in excluded_keys: - # Don't mask keys that are explicitly excluded - masked_data[k] = v - elif self.is_sensitive_key(k): + elif self.is_sensitive_key(k, excluded_keys): str_value = str(v) if v is not None else "" masked_data[k] = self._mask_value(str_value) else: diff --git a/litellm/proxy/common_utils/openai_endpoint_utils.py b/litellm/proxy/common_utils/openai_endpoint_utils.py index 7b1a2945ba6..bedaf31e758 100644 --- a/litellm/proxy/common_utils/openai_endpoint_utils.py +++ b/litellm/proxy/common_utils/openai_endpoint_utils.py @@ -2,7 +2,7 @@ Contains utils used by OpenAI compatible endpoints """ -from typing import Optional +from typing import Optional, Set from fastapi import Request @@ -12,12 +12,16 @@ from litellm.proxy.common_utils.http_parsing_utils import _read_request_body SENSITIVE_DATA_MASKER = SensitiveDataMasker() -def remove_sensitive_info_from_deployment(deployment_dict: dict) -> dict: +def remove_sensitive_info_from_deployment( + deployment_dict: dict, + excluded_keys: Optional[Set[str]] = None, +) -> dict: """ Removes sensitive information from a deployment dictionary. Args: deployment_dict (dict): The deployment dictionary to remove sensitive information from. + excluded_keys (Optional[Set[str]]): Set of keys that should not be masked (exact match). Returns: dict: The modified deployment dictionary with sensitive information removed. @@ -28,7 +32,9 @@ def remove_sensitive_info_from_deployment(deployment_dict: dict) -> dict: deployment_dict["litellm_params"].pop("aws_access_key_id", None) deployment_dict["litellm_params"].pop("aws_secret_access_key", None) - deployment_dict["litellm_params"] = SENSITIVE_DATA_MASKER.mask_dict(deployment_dict["litellm_params"]) + deployment_dict["litellm_params"] = SENSITIVE_DATA_MASKER.mask_dict( + deployment_dict["litellm_params"], excluded_keys=excluded_keys + ) return deployment_dict diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9c98ad215b3..91353e01624 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7069,7 +7069,9 @@ async def model_info_v2( _model["model_info"] = model_info # don't return the api key / vertex credentials # don't return the llm credentials - _model = remove_sensitive_info_from_deployment(_model) + _model = remove_sensitive_info_from_deployment( + _model, excluded_keys={"litellm_credential_name"} + ) verbose_proxy_logger.debug("all_models: %s", all_models) return {"data": all_models} @@ -7536,7 +7538,9 @@ def _get_proxy_model_info(model: dict) -> dict: model_info[k] = v model["model_info"] = model_info # don't return the llm credentials - model = remove_sensitive_info_from_deployment(deployment_dict=model) + model = remove_sensitive_info_from_deployment( + deployment_dict=model, excluded_keys={"litellm_credential_name"} + ) return model @@ -7604,7 +7608,8 @@ async def model_info_v1( # noqa: PLR0915 ) _deployment_info_dict = _deployment_info.model_dump() _deployment_info_dict = remove_sensitive_info_from_deployment( - deployment_dict=_deployment_info_dict + deployment_dict=_deployment_info_dict, + excluded_keys={"litellm_credential_name"}, ) return {"data": _deployment_info_dict} diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py index 6356c5137b8..2836398228a 100644 --- a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py +++ b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py @@ -29,3 +29,49 @@ def test_lists_are_preserved_not_converted_to_strings(): # Must be a list, not a string assert isinstance(masked["tags"], list) assert masked["tags"] == ["East US 2", "production", "test"] + + +def test_excluded_keys_exact_match(): + """ + Test that excluded_keys prevents masking of specific keys (exact match). + """ + masker = SensitiveDataMasker() + + data = { + "api_key": "sk-1234567890abcdef", + "litellm_credentials_name": "my-credential-name", + "access_token": "token-12345", + "port": 6379, + } + + # Without excluded_keys, sensitive keys should be masked + masked = masker.mask_dict(data) + assert masked["api_key"] != "sk-1234567890abcdef" + assert "*" in masked["api_key"] + assert masked["access_token"] != "token-12345" + assert "*" in masked["access_token"] + + # With excluded_keys, litellm_credentials_name should NOT be masked (exact match) + # This ensures that even if pattern matching logic changes, excluded keys won't be masked + masked = masker.mask_dict(data, excluded_keys={"litellm_credentials_name"}) + assert masked["litellm_credentials_name"] == "my-credential-name" + + # Other sensitive keys should still be masked + assert masked["api_key"] != "sk-1234567890abcdef" + assert "*" in masked["api_key"] + assert masked["access_token"] != "token-12345" + assert "*" in masked["access_token"] + + # Non-sensitive keys should remain unchanged + assert masked["port"] == 6379 + + # Test case sensitivity - excluded_keys should be exact match + masked = masker.mask_dict(data, excluded_keys={"LITELLM_CREDENTIALS_NAME"}) + # Should still be masked because case doesn't match (exact match required) + assert masked["litellm_credentials_name"] == "my-credential-name" # Not masked because it doesn't match patterns anyway + + # Test with api_key in excluded_keys to verify it works for keys that would be masked + masked = masker.mask_dict(data, excluded_keys={"api_key"}) + assert masked["api_key"] == "sk-1234567890abcdef" # Should NOT be masked + assert masked["access_token"] != "token-12345" # Should still be masked + assert "*" in masked["access_token"] diff --git a/tests/test_litellm/proxy/common_utils/test_openai_endpoint_utils.py b/tests/test_litellm/proxy/common_utils/test_openai_endpoint_utils.py index a5094c94bd8..a7ce39c2e36 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_endpoint_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_endpoint_utils.py @@ -85,3 +85,37 @@ from litellm.proxy.common_utils.openai_endpoint_utils import remove_sensitive_in def test_remove_sensitive_info_from_deployment(model_config: dict, expected_config: dict): sanitized_config = remove_sensitive_info_from_deployment(model_config) assert sanitized_config == expected_config + + +def test_remove_sensitive_info_from_deployment_with_excluded_keys(): + """ + Test that excluded_keys prevents masking of specific keys (exact match). + """ + model_config = { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "sk-sensitive-key-123", + "litellm_credentials_name": "my-credential-name", + "access_token": "token-12345", + "temperature": 0.7 + } + } + + # Without excluded_keys, access_token should be masked (contains "token") + sanitized_config = remove_sensitive_info_from_deployment(model_config) + assert sanitized_config["litellm_params"]["access_token"] != "token-12345" + assert "*" in sanitized_config["litellm_params"]["access_token"] + + # With excluded_keys, litellm_credentials_name should NOT be masked (even if it would match patterns) + sanitized_config = remove_sensitive_info_from_deployment( + model_config, excluded_keys={"litellm_credentials_name"} + ) + assert sanitized_config["litellm_params"]["litellm_credentials_name"] == "my-credential-name" + + # access_token should still be masked (not in excluded_keys) + assert sanitized_config["litellm_params"]["access_token"] != "token-12345" + assert "*" in sanitized_config["litellm_params"]["access_token"] + + # api_key should still be removed (popped) regardless of excluded_keys + assert "api_key" not in sanitized_config["litellm_params"]