From 02131cfede7ae203b0afcb0cdadf945358a35a99 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:32:43 -0700 Subject: [PATCH 1/7] feat(proxy): encrypt guardrail params, nested litellm_params and router_settings secrets at rest --- .../common_utils/encrypt_decrypt_utils.py | 89 ++++++++- .../common_utils/registry_read_through.py | 4 +- litellm/proxy/db/master_key_migration.py | 1 + .../proxy/guardrails/guardrail_endpoints.py | 9 +- .../proxy/guardrails/guardrail_registry.py | 31 ++- litellm/proxy/guardrails/usage_endpoints.py | 12 +- .../credential_migration.py | 1 + .../fallback_management_endpoints.py | 5 +- .../key_management_endpoints.py | 38 ++++ .../model_management_endpoints.py | 64 +++++-- litellm/proxy/proxy_server.py | 24 ++- litellm/repositories/model_repository.py | 28 +-- .../test_encrypt_decrypt_utils.py | 128 ++++++++++++- .../test_registry_read_through.py | 39 ++++ .../proxy/db/test_master_key_migration.py | 36 +++- .../guardrails/test_guardrail_endpoints.py | 113 +++++++++++ .../guardrails/test_guardrail_registry.py | 118 +++++++++++- .../proxy/guardrails/test_usage_endpoints.py | 22 +++ .../test_credential_migration.py | 21 +++ .../test_key_management_endpoints.py | 63 +++++++ .../test_model_management_endpoints.py | 178 +++++++++++++++--- .../proxy/proxy_server/test_proxy_config.py | 19 ++ .../test_fallback_management_endpoints.py | 33 ++++ tests/test_litellm/proxy/test_proxy_server.py | 53 ++++++ tests/unit/repositories/test_repositories.py | 110 +++++------ 25 files changed, 1088 insertions(+), 151 deletions(-) diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 3584aaaf833..b83aa287ba9 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -3,9 +3,11 @@ import os from collections.abc import Mapping from typing import Final, Literal, cast -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # Versioned ciphertext marker for AES-256-GCM values. # Format: "v2:gcm:" + base64url(nonce(12) || ciphertext || tag(16)). @@ -258,3 +260,88 @@ def decode_secret_map(value: object, *, key: str) -> Mapping[str, str] | None: return _SECRET_MAP.validate_json(decrypted, strict=True) except ValidationError: raise SecretMapDecodeError(f"Cannot decode encrypted MCP {key}; check LITELLM_SALT_KEY") from None + + +_JSON_VALUE: Final = TypeAdapter(JsonValue) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def json_value(value: object) -> JsonValue: + plain: Final = dict(value) if isinstance(value, Mapping) else value + try: + return _JSON_VALUE.validate_python(plain) + except ValidationError: + return _JSON_VALUE.validate_json(safe_dumps(plain)) + + +def _plaintext_of(value: str) -> str: + signing_key: Final = _get_salt_key() + if signing_key is None: + return value + decrypted: Final = decrypt_if_encrypted_with(value, signing_key) + return value if decrypted is None else decrypted + + +def _encrypted_string(value: str, new_encryption_key: str | None) -> str: + if (new_encryption_key or _get_salt_key()) is None: + return value + return _SECRET_STRING.validate_python( + encrypt_value_helper(_plaintext_of(value), new_encryption_key=new_encryption_key), strict=True + ) + + +def encrypt_json_strings(value: JsonValue, new_encryption_key: str | None = None, depth: int = 0) -> JsonValue: + if depth > DEFAULT_MAX_RECURSE_DEPTH: + return value + match value: + case str(): + return _encrypted_string(value, new_encryption_key) + case list(): + return [encrypt_json_strings(item, new_encryption_key, depth + 1) for item in value] + case dict(): + return {key: encrypt_json_strings(item, new_encryption_key, depth + 1) for key, item in value.items()} + case _: + return value + + +def decrypt_json_strings(value: JsonValue, depth: int = 0) -> JsonValue: + if depth > DEFAULT_MAX_RECURSE_DEPTH: + return value + match value: + case str(): + return _plaintext_of(value) + case list(): + return [decrypt_json_strings(item, depth + 1) for item in value] + case dict(): + return {key: decrypt_json_strings(item, depth + 1) for key, item in value.items()} + case _: + return value + + +def _stored_json_object(value: object) -> dict[str, JsonValue]: + try: + return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value) + except ValidationError: + return {} + + +def decrypt_stored_json_object(value: object) -> dict[str, JsonValue]: + return {key: decrypt_json_strings(item) for key, item in _stored_json_object(value).items()} + + +ENCRYPTED_CONFIG_SECTIONS: Final = frozenset({"router_settings"}) + + +def encrypt_config_section( + section_name: str, section: object, new_encryption_key: str | None = None +) -> dict[str, JsonValue]: + values: Final = _stored_json_object(section) + if section_name not in ENCRYPTED_CONFIG_SECTIONS: + return values + return {key: encrypt_json_strings(item, new_encryption_key=new_encryption_key) for key, item in values.items()} + + +def decrypt_config_section(section_name: str, value: object) -> dict[str, JsonValue]: + if section_name in ENCRYPTED_CONFIG_SECTIONS: + return decrypt_stored_json_object(value) + return _stored_json_object(value) diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 63da5d15207..6723235db9f 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -136,9 +136,9 @@ async def _resync_guardrails(guardrail_name: str) -> bool: from litellm.proxy.guardrails.guardrail_registry import ( GUARDRAIL_RECONCILE_LOCK, IN_MEMORY_GUARDRAIL_HANDLER, + guardrail_from_row, ) from litellm.repositories.table_repositories import GuardrailsRepository - from litellm.types.guardrails import Guardrail if not _db_backed_registries_enabled("guardrails"): return False @@ -152,7 +152,7 @@ async def _resync_guardrails(guardrail_name: str) -> bool: if row is None: return False async with GUARDRAIL_RECONCILE_LOCK: - IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=Guardrail(**dict(row))) + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=guardrail_from_row(row)) return _initialized_guardrail(guardrail_name) is not None diff --git a/litellm/proxy/db/master_key_migration.py b/litellm/proxy/db/master_key_migration.py index d100554201a..5ac44e4e791 100644 --- a/litellm/proxy/db/master_key_migration.py +++ b/litellm/proxy/db/master_key_migration.py @@ -27,6 +27,7 @@ _SECRET_COLUMNS: Final = ( _SecretColumn("LiteLLM_ProxyModelTable", "model_id", "litellm_params"), _SecretColumn("LiteLLM_CredentialsTable", "credential_id", "credential_values"), _SecretColumn("LiteLLM_Config", "param_name", "param_value"), + _SecretColumn("LiteLLM_GuardrailsTable", "guardrail_id", "litellm_params"), _SecretColumn("LiteLLM_SSOConfig", "id", "sso_settings"), _SecretColumn("LiteLLM_CacheConfig", "id", "cache_settings"), _SecretColumn("LiteLLM_ConfigOverrides", "config_type", "config_value"), diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 6874d7aa73e..ec4a6a59557 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -21,12 +21,13 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_stored_json_object from litellm.proxy.common_utils.path_utils import safe_join from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( build_sandbox_globals, compile_sandboxed, ) -from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry +from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry, encrypted_guardrail_litellm_params from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.repositories.prisma_protocols import TableActions @@ -765,7 +766,7 @@ async def register_guardrail( raise HTTPException(status_code=500, detail=str(e)) now: Final = datetime.now(timezone.utc) - litellm_params_str: Final = safe_dumps(params) + litellm_params_str: Final = encrypted_guardrail_litellm_params(params) guardrail_info: Final = dict(request.guardrail_info or {}) guardrail_info["submitted_by_user_id"] = user_api_key_dict.user_id guardrail_info["submitted_by_email"] = user_api_key_dict.user_email @@ -839,7 +840,7 @@ def _row_to_submission_item(row: "LiteLLM_GuardrailsTable") -> GuardrailSubmissi guardrail_info: Final = _parse_json_field(row.guardrail_info) or {} team_guardrail: Final = row.team_id is not None - raw_params: Final = _parse_json_field(row.litellm_params) or {} + raw_params: Final = decrypt_stored_json_object(_parse_json_field(row.litellm_params) or {}) masked_params: Final = _get_masked_values(raw_params, unmasked_length=4, number_of_asterisks=4) return GuardrailSubmissionItem( guardrail_id=row.guardrail_id, @@ -1034,7 +1035,7 @@ async def approve_guardrail_submission( guardrail_dict: Final = { "guardrail_id": row.guardrail_id, "guardrail_name": row.guardrail_name, - "litellm_params": litellm_params, + "litellm_params": decrypt_stored_json_object(litellm_params), "guardrail_info": guardrail_info or {}, "team_id": row.team_id, } diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 0dc50cd6196..5ddc9238d1c 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -2,6 +2,7 @@ import asyncio import importlib +import json import os from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime, timezone @@ -20,6 +21,11 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_tool_message_for_guardrail, ) +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_stored_json_object, + encrypt_json_strings, + json_value, +) from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, ) @@ -77,6 +83,18 @@ def _guardrail_table(prisma_client: PrismaClient) -> "TableActions[prisma_models return GuardrailsRepository(prisma_client).table +def encrypted_guardrail_litellm_params(litellm_params: Mapping[str, object]) -> str: + return json.dumps(encrypt_json_strings(json_value(litellm_params))) + + +def guardrail_from_row(row: _GuardrailRowLike) -> Guardrail: + fields: Final = dict(row) + stored_params: Final = fields.get("litellm_params") + if stored_params is None: + return Guardrail(**fields) + return Guardrail(**{**fields, "litellm_params": decrypt_stored_json_object(stored_params)}) + + guardrail_initializer_registry: Final = { SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock, SupportedGuardrailIntegrations.LAKERA.value: initialize_lakera, @@ -295,7 +313,7 @@ class GuardrailRegistry: litellm_params_dict = litellm_params_obj.model_dump() else: litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} - litellm_params: Final[str] = safe_dumps(litellm_params_dict) + litellm_params: Final[str] = encrypted_guardrail_litellm_params(litellm_params_dict) guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {})) # Create guardrail in DB @@ -341,7 +359,7 @@ class GuardrailRegistry: litellm_params_dict = litellm_params_obj.model_dump() else: litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} - litellm_params: Final[str] = safe_dumps(litellm_params_dict) + litellm_params: Final[str] = encrypted_guardrail_litellm_params(litellm_params_dict) guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {})) # Update in DB @@ -357,8 +375,7 @@ class GuardrailRegistry: if updated_guardrail is None: raise ValueError(f"Guardrail not found, passed guardrail_id={guardrail_id}") - # Convert to dict and return - return dict(updated_guardrail) + return dict(guardrail_from_row(updated_guardrail)) except Exception as e: raise Exception(f"Error updating guardrail in DB: {e}") @@ -378,7 +395,7 @@ class GuardrailRegistry: guardrails: Final[list[Guardrail]] = [] for guardrail in guardrails_from_db: - guardrails.append(Guardrail(**(dict(guardrail)))) + guardrails.append(guardrail_from_row(guardrail)) return guardrails except Exception as e: @@ -394,7 +411,7 @@ class GuardrailRegistry: if not guardrail: return None - return Guardrail(**(dict(guardrail))) + return guardrail_from_row(guardrail) except Exception as e: raise Exception(f"Error getting guardrail from DB: {e}") @@ -410,7 +427,7 @@ class GuardrailRegistry: if not guardrail: return None - return Guardrail(**(dict(guardrail))) + return guardrail_from_row(guardrail) except Exception as e: raise Exception(f"Error getting guardrail from DB: {e}") diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index fc237bd55c4..21cfe382fad 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -17,6 +17,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_stored_json_object from litellm.proxy.guardrails.usage_tracking import guardrail_status_to_action from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( @@ -429,6 +430,13 @@ def _to_dict(value: object) -> dict[str, Any]: return {} +def _guardrail_litellm_params(g: "_DbOrConfigGuardrail") -> Mapping[str, object]: + value: Final = _get_guardrail_field(g, "litellm_params") + if isinstance(value, dict): + return decrypt_stored_json_object(value) + return _to_dict(value) + + def _field_str(mapping: Mapping[str, object], key: str, default: str) -> str: """Stringify `mapping[key]`, falling back to `default` when the key is absent.""" return str(mapping.get(key, default)) @@ -463,7 +471,7 @@ def _guardrail_overview_rows( break req, blocked = a["requests"], a["blocked"] fail_rate = (100.0 * blocked / req) if req else 0.0 - litellm_params = _to_dict(_get_guardrail_field(g, "litellm_params")) + litellm_params = _guardrail_litellm_params(g) provider = _field_str(litellm_params, "guardrail", "Unknown") guardrail_info = _to_dict(_get_guardrail_field(g, "guardrail_info")) gtype = _field_str(guardrail_info, "type", "Guardrail") @@ -703,7 +711,7 @@ async def guardrails_usage_detail( {"date": d, "passed": v["passed"], "blocked": v["blocked"], "score": None} for d, v in sorted(ts_by_date.items()) ] - litellm_params: Final = _to_dict(_get_guardrail_field(guardrail, "litellm_params")) + litellm_params: Final = _guardrail_litellm_params(guardrail) guardrail_info: Final = _to_dict(_get_guardrail_field(guardrail, "guardrail_info")) _guardrail_name: Final = _get_guardrail_field(guardrail, "guardrail_name") daily_unit_sums: Final = sorted(_by(units_rows, lambda r: r.date, _sum_counter_units).items()) diff --git a/litellm/proxy/management_endpoints/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py index 915cce87dbd..7ab8966549c 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -457,6 +457,7 @@ def _classify_callback_value(value: object) -> ValueClass: _COVERED_TABLE_SPECS: Final = [ ("model_table", "litellm_proxymodeltable", ("litellm_params",), ()), ("credentials", "litellm_credentialstable", ("credential_values",), ()), + ("guardrails", "litellm_guardrailstable", ("litellm_params",), ()), ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()), ("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)), ("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)), diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index 543538f5edd..3ca86e27575 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -17,6 +17,7 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.model_checks import get_all_fallbacks from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_config_section if TYPE_CHECKING: from fastapi import APIRouter, Depends, HTTPException, status @@ -153,7 +154,7 @@ async def create_fallback( router_settings[fallback_key] = existing_fallbacks # Save to database - convert router_settings to JSON string - router_settings_json: Final = json.dumps(router_settings) + router_settings_json: Final = json.dumps(encrypt_config_section("router_settings", router_settings)) await ConfigRepository(prisma_client).table.upsert( where={"param_name": "router_settings"}, data={ @@ -324,7 +325,7 @@ async def delete_fallback( router_settings[fallback_key] = updated_fallbacks # Save to database - convert router_settings to JSON string - router_settings_json: Final = json.dumps(router_settings) + router_settings_json: Final = json.dumps(encrypt_config_section("router_settings", router_settings)) await ConfigRepository(prisma_client).table.upsert( where={"param_name": "router_settings"}, data={ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7e159ec90e7..5c4c369dd73 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -80,6 +80,12 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, publish_config_change, ) +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + ENCRYPTED_CONFIG_SECTIONS, + encrypt_config_section, + encrypt_json_strings, + json_value, +) from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -142,6 +148,7 @@ from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( DeletedVerificationTokenRepository, DeprecatedVerificationTokenRepository, + GuardrailsRepository, ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository @@ -5216,6 +5223,13 @@ async def delete_key_aliases( ) +async def _guardrail_rows(prisma_client: PrismaClient) -> "Sequence[prisma_models.LiteLLM_GuardrailsTable]": + try: + return await GuardrailsRepository(prisma_client).table.find_many() + except Exception: # noqa: BLE001 # a deployment without the guardrails table still rotates every other store + return () + + async def _rotate_master_key( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, @@ -5301,6 +5315,30 @@ async def _rotate_master_key( data={"param_value": prisma.Json(encrypted_env_vars)}, ) + for c in config: + if c.param_name in ENCRYPTED_CONFIG_SECTIONS and c.param_value is not None: + await _config_table(prisma_client).update( + where={"param_name": c.param_name}, + data={ + "param_value": prisma.Json( + encrypt_config_section(c.param_name, c.param_value, new_encryption_key=new_master_key) + ) + }, + ) + + # 3b. process guardrails table + for guardrail_row in await _guardrail_rows(prisma_client): + if guardrail_row.litellm_params is None: + continue + await GuardrailsRepository(prisma_client).table.update( + where={"guardrail_id": guardrail_row.guardrail_id}, + data={ + "litellm_params": prisma.Json( + encrypt_json_strings(json_value(guardrail_row.litellm_params), new_encryption_key=new_master_key) + ) + }, + ) + # 4. process MCP server table try: await rotate_mcp_server_credentials_master_key( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ae294871afc..e9d85060599 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -22,7 +22,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator import litellm from litellm._logging import verbose_proxy_logger @@ -64,7 +64,8 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( ) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, - encrypt_value_helper, + encrypt_json_strings, + json_value, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient @@ -436,6 +437,33 @@ def _effective_complexity_router_config( } +def _encrypted_complexity_router_config(config: object, new_encryption_key: str | None = None) -> JsonValue: + value: Final = json_value(config) + if not isinstance(value, dict): + return value + jev: Final = value.get("jev_classifier_config") + if not isinstance(jev, dict): + return value + api_key: Final = jev.get("api_key") + if not isinstance(api_key, str): + return value + encrypted_api_key: Final = encrypt_json_strings(api_key, new_encryption_key=new_encryption_key) + return {**value, "jev_classifier_config": {**jev, "api_key": encrypted_api_key}} + + +def _encrypted_litellm_params( + litellm_params: Mapping[str, object], new_encryption_key: str | None = None +) -> dict[str, JsonValue]: + return { + key: ( + _encrypted_complexity_router_config(value, new_encryption_key) + if key == "complexity_router_config" + else encrypt_json_strings(json_value(value), new_encryption_key=new_encryption_key) + ) + for key, value in litellm_params.items() + } + + def _effective_model( incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None ) -> str | None: @@ -986,19 +1014,16 @@ def update_db_model( k: v for k, v in stored_model_info.items() if k not in echoed_pricing } - # update litellm params if updated_patch.litellm_params: - # Encrypt any sensitive values - encrypted_params: Final = { + patched_params: Final = { k: ( _effective_complexity_router_config(updated_patch.litellm_params, db_model.litellm_params) if k == "complexity_router_config" - else encrypt_value_helper(v) + else v ) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() } - - merged_litellm_params.update(encrypted_params) + merged_litellm_params.update(_encrypted_litellm_params({**merged_litellm_params, **patched_params})) # update model info if updated_patch.model_info: @@ -1502,9 +1527,8 @@ async def _add_model_to_db( # encrypt litellm params # _litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) _original_litellm_model_name: Final = model_params.litellm_params.model - for k, v in _litellm_params_dict.items(): - encrypted_value = encrypt_value_helper(value=v, new_encryption_key=new_encryption_key) - model_params.litellm_params[k] = encrypted_value + for k, v in _encrypted_litellm_params(_litellm_params_dict, new_encryption_key).items(): + model_params.litellm_params[k] = v _data: Final[dict] = { "model_id": model_params.model_info.id, "model_name": model_params.model_name, @@ -2659,25 +2683,25 @@ async def update_model( _new_litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) - ### ENCRYPT PARAMS ### - encrypted_params: Final = MappingProxyType( + requested_params: Final = MappingProxyType( { k: ( _effective_complexity_router_config(model_params.litellm_params, deployment.litellm_params) if k == "complexity_router_config" - else encrypt_value_helper(value=v) + else v ) for k, v in _new_litellm_params_dict.items() } ) - ### MERGE WITH EXISTING DATA ### _mp: Final[dict[str, object]] = model_params.litellm_params.dict() - merged_dictionary: Final = { - key: _existing_litellm_params_dict[key] if value is None else encrypted_params[key] - for key, value in _mp.items() - if value is not None or _existing_litellm_params_dict.get(key) is not None - } + merged_dictionary: Final = _encrypted_litellm_params( + { + key: _existing_litellm_params_dict[key] if value is None else requested_params[key] + for key, value in _mp.items() + if value is not None or _existing_litellm_params_dict.get(key) is not None + } + ) renamed_to: Final = ( model_params.model_name diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61304f0d919..98c571cc6dd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -404,8 +404,13 @@ from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.discoverable_model_filter import discoverable_rows, undiscoverable_model_names from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + ENCRYPTED_CONFIG_SECTIONS, + decrypt_config_section, + decrypt_json_strings, decrypt_value_helper, + encrypt_config_section, encrypt_value_helper, + json_value, ) from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.healthy_model_filter import ( @@ -5382,10 +5387,8 @@ class ProxyConfig: existing_row: Final[_ConfigParamRow | None] = await config_table.find_first(where=config_where) existing_value: Final[object] = cast(object, existing_row.param_value) if existing_row is not None else None existing_section: Final[Mapping[str, JsonValue]] = ( - _CONFIG_SECTION_VALUES.validate_json(existing_value) - if isinstance(existing_value, str) - else _CONFIG_SECTION_VALUES.validate_python(existing_value) - if isinstance(existing_value, Mapping) + decrypt_config_section(section_name, existing_value) + if isinstance(existing_value, (str, Mapping)) else MappingProxyType({}) ) merged_section: Final[Mapping[str, JsonValue]] = MappingProxyType( @@ -5399,7 +5402,7 @@ class ProxyConfig: ) if merged_section == existing_section: return None - serialized_section: Final = json.dumps(dict(merged_section)) # mutable-ok: JSON encoder requires a dict + serialized_section: Final = json.dumps(encrypt_config_section(section_name, merged_section)) config_data: Final[_ConfigParamUpsert] = { "create": {"param_name": section_name, "param_value": serialized_section}, "update": {"param_value": serialized_section}, @@ -7051,6 +7054,8 @@ class ProxyConfig: return frozenset(combined_id_list) | kept_config_ids def _resolve_db_litellm_param(self, key: str, value: object) -> object: + if isinstance(value, (dict, list)): + return decrypt_json_strings(json_value(value)) if not isinstance(value, str): return value @@ -7419,7 +7424,7 @@ class ProxyConfig: where={"param_name": "router_settings"} ) db_values: Final = ( - _as_settings_mapping(db_router_settings.param_value) + _as_settings_mapping(decrypt_config_section("router_settings", db_router_settings.param_value)) if db_router_settings is not None and db_router_settings.param_value is not None else _EMPTY_SETTINGS_MAPPING ) @@ -7715,7 +7720,8 @@ class ProxyConfig: os.environ[key] = decrypted_value return _as_settings_mapping(normalized) - scrubbed: Final = _scrub_db_overlay_remote_module_loads(section=section, db_value=value) + db_value: Final = decrypt_config_section(section, value) if section in ENCRYPTED_CONFIG_SECTIONS else value + scrubbed: Final = _scrub_db_overlay_remote_module_loads(section=section, db_value=db_value) return _as_settings_mapping(scrubbed) def _apply_litellm_settings_db_values(self, db_values: Mapping[str, SettingsJsonValue]) -> None: @@ -17744,10 +17750,10 @@ async def update_config( ) if row is None or row.param_value is None: return {} - return dict(row.param_value) + return decrypt_config_section(param_name, row.param_value) async def _upsert_section(param_name: str, value: dict) -> None: - serialized: Final = json.dumps(value) + serialized: Final = json.dumps(encrypt_config_section(param_name, value)) await ConfigRepository(prisma_client).table.upsert( where={"param_name": param_name}, data={ diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 8ee76b93923..c21c8c4f653 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -8,8 +8,9 @@ from typing import TYPE_CHECKING, Any, Final from litellm.models.model import LiteLLM_ProxyModelTable from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, - encrypt_value_helper, + decrypt_json_strings, + encrypt_json_strings, + json_value, ) from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions @@ -39,26 +40,13 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): return LiteLLM_ProxyModelTable def _encrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]: - """Encrypt sensitive values in litellm_params.""" - encrypted: Final = {} - for key, value in litellm_params.items(): - if isinstance(value, str): - encrypted[key] = encrypt_value_helper(value, new_encryption_key=self._encryption_key) - else: - encrypted[key] = value - return encrypted + return { + key: encrypt_json_strings(json_value(value), new_encryption_key=self._encryption_key) + for key, value in litellm_params.items() + } def _decrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]: - """Decrypt sensitive values in litellm_params.""" - decrypted: Final = {} - for key, value in litellm_params.items(): - if isinstance(value, str): - decrypted[key] = decrypt_value_helper( - value, key=key, exception_type="debug", return_original_value=True - ) - else: - decrypted[key] = value - return decrypted + return {key: decrypt_json_strings(json_value(value)) for key, value in litellm_params.items()} def _to_model(self, record: Any) -> LiteLLM_ProxyModelTable | None: """Convert a database record to a Model with decryption.""" diff --git a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py index 9c07242bd23..d32f1c37a8b 100644 --- a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -7,18 +7,29 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130 """ import base64 +import json +from types import MappingProxyType import pytest +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import ( _V2_GCM_PREFIX, + decrypt_config_section, decrypt_if_encrypted_with, + decrypt_json_strings, + decrypt_stored_json_object, decrypt_value_helper, + encrypt_config_section, + encrypt_json_strings, encrypt_value, encrypt_value_helper, + json_value, ) +SALT_KEY = "sk-salt-aes-1234" + def _use_aes(monkeypatch): """Flip the write-time algorithm to AES-256-GCM for the duration of a test.""" @@ -28,7 +39,7 @@ def _use_aes(monkeypatch): @pytest.fixture(autouse=True) def _salt_key(monkeypatch): # Dominant convention in the test_litellm/ tree: set the key via env. - monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-aes-1234") + monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY) # Ensure the legacy default is in force unless a test opts into AES. monkeypatch.setattr(proxy_server, "general_settings", {}) yield @@ -236,3 +247,118 @@ def test_explicit_key_decrypt_supports_the_empty_master_key(): written_with_empty_key = encrypt_value(value="stored-secret", signing_key="") assert decrypt_if_encrypted_with(base64.urlsafe_b64encode(written_with_empty_key).decode(), "") == "stored-secret" + + +_NESTED_PARAMS = { + "model": "openai/gpt-5.4-mini", + "extra_headers": {"Authorization": "Bearer gateway-secret", "X-Trace": "trace-1"}, + "fallbacks": [{"gpt-5.5-mini": ["gpt-5.4-mini"]}], + "rpm": 10, + "enabled": True, + "api_base": None, +} + + +def _nested_string_leaves(stored: dict) -> list: + return [ + stored["model"], + stored["extra_headers"]["Authorization"], + stored["extra_headers"]["X-Trace"], + stored["fallbacks"][0]["gpt-5.5-mini"][0], + ] + + +@pytest.mark.parametrize("use_aes", [False, True]) +def test_encrypt_json_strings_encrypts_every_nested_string_leaf(monkeypatch, use_aes: bool): + if use_aes: + _use_aes(monkeypatch) + + stored = encrypt_json_strings(_NESTED_PARAMS) + + assert (stored["rpm"], stored["enabled"], stored["api_base"]) == (10, True, None) + plaintexts = _nested_string_leaves(_NESTED_PARAMS) + assert all(leaf != plain for leaf, plain in zip(_nested_string_leaves(stored), plaintexts)) + assert [decrypt_if_encrypted_with(leaf, SALT_KEY) for leaf in _nested_string_leaves(stored)] == plaintexts + assert decrypt_json_strings(stored) == _NESTED_PARAMS + + +def test_encrypt_json_strings_reencrypts_ciphertext_under_the_new_key_without_double_wrapping(): + stored = encrypt_json_strings({"api_key": "top-secret", "nested": {"token": encrypt_value_helper("nested-secret")}}) + + rotated = encrypt_json_strings(stored, new_encryption_key="sk-rotated") + + assert decrypt_if_encrypted_with(rotated["api_key"], "sk-rotated") == "top-secret" + assert decrypt_if_encrypted_with(rotated["nested"]["token"], "sk-rotated") == "nested-secret" + assert decrypt_if_encrypted_with(rotated["api_key"], SALT_KEY) is None + + +def test_encrypt_json_strings_treats_a_legacy_plaintext_row_as_plaintext(): + legacy_row = {"api_key": "plain-secret", "extra_headers": {"Authorization": "Bearer plain"}} + + assert decrypt_json_strings(legacy_row) == legacy_row + assert decrypt_json_strings(encrypt_json_strings(legacy_row)) == legacy_row + + +def test_encrypt_json_strings_without_any_key_stores_the_value_as_sent(monkeypatch): + monkeypatch.delenv("LITELLM_SALT_KEY") + monkeypatch.setattr(proxy_server, "master_key", None) + + assert encrypt_json_strings(_NESTED_PARAMS) == _NESTED_PARAMS + assert decrypt_json_strings(_NESTED_PARAMS) == _NESTED_PARAMS + + +def test_encrypt_json_strings_stops_descending_past_the_recursion_cap(): + leaf = {"secret": "deep-secret"} + value = leaf + for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1): + value = {"child": value} + + stored = encrypt_json_strings(value) + + deepest = stored + for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1): + deepest = deepest["child"] + assert deepest == leaf + + +def test_json_value_coerces_non_json_containers(): + assert json_value(MappingProxyType({"headers": ("a", "b")})) == {"headers": ["a", "b"]} + + +@pytest.mark.parametrize("stored", ["not json", "[1]", ["a"], None, 3]) +def test_decrypt_stored_json_object_returns_an_empty_object_for_a_non_object_row(stored: object): + assert decrypt_stored_json_object(stored) == {} + + +def test_decrypt_stored_json_object_reads_json_text_and_parsed_rows_alike(): + stored = encrypt_json_strings(_NESTED_PARAMS) + + assert decrypt_stored_json_object(json.dumps(stored)) == _NESTED_PARAMS + assert decrypt_stored_json_object(stored) == _NESTED_PARAMS + + +def test_encrypt_config_section_encrypts_router_settings_string_leaves_only(): + router_settings = {"redis_password": "redis-pw", "num_retries": 2, "fallbacks": [{"gpt-5.5-mini": ["gpt-5.4-mini"]}]} + + stored = encrypt_config_section("router_settings", router_settings) + + assert stored["num_retries"] == 2 + assert stored["redis_password"] != "redis-pw" + assert decrypt_if_encrypted_with(stored["redis_password"], SALT_KEY) == "redis-pw" + assert decrypt_if_encrypted_with(stored["fallbacks"][0]["gpt-5.5-mini"][0], SALT_KEY) == "gpt-5.4-mini" + assert decrypt_config_section("router_settings", json.dumps(stored)) == router_settings + + +def test_encrypt_config_section_leaves_other_sections_readable(): + general_settings = {"alerting": ["slack"], "proxy_batch_write_at": 60} + + assert encrypt_config_section("general_settings", general_settings) == general_settings + assert decrypt_config_section("general_settings", json.dumps(general_settings)) == general_settings + + +def test_encrypt_config_section_rekeys_router_settings_for_master_key_rotation(): + stored = encrypt_config_section("router_settings", {"redis_password": "redis-pw"}) + + rotated = encrypt_config_section("router_settings", stored, new_encryption_key="sk-rotated") + + assert decrypt_if_encrypted_with(rotated["redis_password"], "sk-rotated") == "redis-pw" diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index ca2ff8bcce1..fb2a2b742f9 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -300,6 +300,45 @@ async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sib IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id) +class EncryptedGuardrailRow(FakeGuardrailRow): + def __iter__(self): + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_json_strings + + plaintext_row: Final = dict(super().__iter__()) + return iter({**plaintext_row, "litellm_params": encrypt_json_strings(plaintext_row["litellm_params"])}.items()) + + +@pytest.mark.asyncio +async def test_get_guardrail_with_read_through_decrypts_an_encrypted_db_row(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter import ContentFilterGuardrail + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-read-through-salt-1234") + monkeypatch.setattr(proxy_server, "general_settings", {}) + guardrail_id: Final = "read-through-encrypted-guardrail-id" + guardrail_name: Final = "read-through-encrypted-guardrail" + row: Final = EncryptedGuardrailRow(guardrail_id, guardrail_name) + assert dict(row)["litellm_params"]["blocked_words"][0]["keyword"] != "secret" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + try: + guardrail: Final = await get_initialized_guardrail_with_read_through(guardrail_name=guardrail_name) + assert isinstance(guardrail, ContentFilterGuardrail) + assert guardrail.guardrail_name == guardrail_name + assert guardrail._check_blocked_words("this holds a SECRET word") is not None + finally: + IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id) + + @pytest.mark.asyncio async def test_get_guardrail_with_read_through_returns_none_for_unknown_guardrail(monkeypatch): from unittest.mock import AsyncMock, MagicMock diff --git a/tests/test_litellm/proxy/db/test_master_key_migration.py b/tests/test_litellm/proxy/db/test_master_key_migration.py index 9c0fc163b9f..240438f1a6d 100644 --- a/tests/test_litellm/proxy/db/test_master_key_migration.py +++ b/tests/test_litellm/proxy/db/test_master_key_migration.py @@ -98,10 +98,22 @@ def _seeded_tables() -> Tables: "rpm": 10, "use_in_pass_through": False, "api_base": None, + "extra_headers": {"Authorization": _encrypted("Bearer nested-secret"), "X-Trace": "plain"}, }, }, {"model_id": "model-2", "litellm_params": {"api_key": _encrypted("other-deployment", UNRELATED_KEY)}}, ], + "LiteLLM_GuardrailsTable": [ + { + "guardrail_id": "guardrail-1", + "litellm_params": { + "guardrail": _encrypted("openai_moderation"), + "api_key": _encrypted("guardrail-key"), + "default_on": False, + }, + }, + {"guardrail_id": "guardrail-plain", "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}}, + ], "LiteLLM_Config": [ {"param_name": "environment_variables", "param_value": {"LANGFUSE_SECRET_KEY": _encrypted("env-secret")}}, {"param_name": "general_settings", "param_value": {"proxy_batch_write_at": 10, "ui_name": "plain text"}}, @@ -140,14 +152,19 @@ def _seeded_tables() -> Tables: } -_VALUES_UNDER_THE_PREVIOUS_KEY = 8 +_VALUES_UNDER_THE_PREVIOUS_KEY = 11 @pytest.mark.asyncio async def test_reencryption_moves_every_stored_shape_to_the_new_key_and_nothing_else(): tables = _seeded_tables() untouched_before = json.dumps( - [tables["LiteLLM_ProxyModelTable"][1], tables["LiteLLM_Config"][1:], tables["LiteLLM_TeamTable"][1]] + [ + tables["LiteLLM_ProxyModelTable"][1], + tables["LiteLLM_Config"][1:], + tables["LiteLLM_TeamTable"][1], + tables["LiteLLM_GuardrailsTable"][1], + ] ) migrated = await reencrypt_stored_values(_FakeDatabase(tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY) @@ -159,6 +176,13 @@ async def test_reencryption_moves_every_stored_shape_to_the_new_key_and_nothing_ assert decrypt_if_encrypted_with(model_params["model"], NEW_KEY) == "openai/gpt-5.4-mini" assert decrypt_if_encrypted_with(model_params["api_key"], PREVIOUS_KEY) is None assert (model_params["rpm"], model_params["use_in_pass_through"], model_params["api_base"]) == (10, False, None) + assert decrypt_if_encrypted_with(model_params["extra_headers"]["Authorization"], NEW_KEY) == "Bearer nested-secret" + assert model_params["extra_headers"]["X-Trace"] == "plain" + guardrail_params = tables["LiteLLM_GuardrailsTable"][0]["litellm_params"] + assert isinstance(guardrail_params, dict) + assert decrypt_if_encrypted_with(guardrail_params["api_key"], NEW_KEY) == "guardrail-key" + assert decrypt_if_encrypted_with(guardrail_params["guardrail"], NEW_KEY) == "openai_moderation" + assert guardrail_params["default_on"] is False mcp_server = tables["LiteLLM_MCPServerTable"][0] assert decrypt_if_encrypted_with(mcp_server["static_headers"], NEW_KEY) == '{"X-Api-Key": "header-secret"}' assert mcp_server["credentials"]["aws_region_name"] == "us-east-1" @@ -171,7 +195,12 @@ async def test_reencryption_moves_every_stored_shape_to_the_new_key_and_nothing_ assert callback_secret.startswith("litellm_enc::") assert decrypt_if_encrypted_with(callback_secret.removeprefix("litellm_enc::"), NEW_KEY) == "team-callback-secret" assert untouched_before == json.dumps( - [tables["LiteLLM_ProxyModelTable"][1], tables["LiteLLM_Config"][1:], tables["LiteLLM_TeamTable"][1]] + [ + tables["LiteLLM_ProxyModelTable"][1], + tables["LiteLLM_Config"][1:], + tables["LiteLLM_TeamTable"][1], + tables["LiteLLM_GuardrailsTable"][1], + ] ) @@ -197,6 +226,7 @@ async def test_only_rows_holding_values_under_the_previous_key_are_written(): assert sorted(database.writes) == [ ("LiteLLM_Config", "param_value", "environment_variables"), + ("LiteLLM_GuardrailsTable", "litellm_params", "guardrail-1"), ("LiteLLM_MCPServerTable", "credentials", "mcp-1"), ("LiteLLM_MCPServerTable", "env_vars", "mcp-1"), ("LiteLLM_MCPServerTable", "static_headers", "mcp-1"), diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index bf641fd6cd0..7d133558f83 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -2725,3 +2725,116 @@ def test_field_type_inference_handles_pep604_unions(): assert _get_field_type_from_annotation(list[str] | None) == "array" assert _get_field_type_from_annotation(bool | None) == "boolean" assert _unwrap_optional_type(str | None) is str + + +ENDPOINT_SALT_KEY = "sk-endpoint-salt-1234" + + +@pytest.fixture +def endpoint_salt_key(monkeypatch): + from litellm.proxy import proxy_server + + monkeypatch.setenv("LITELLM_SALT_KEY", ENDPOINT_SALT_KEY) + monkeypatch.setattr(proxy_server, "general_settings", {}) + return ENDPOINT_SALT_KEY + + +@pytest.mark.asyncio +async def test_register_guardrail_stores_litellm_params_as_ciphertext(mocker, endpoint_salt_key): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with + + mock_prisma = mocker.Mock() + mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_guardrailstable.create = AsyncMock( + return_value=mocker.Mock( + guardrail_id="reg-encrypted", + guardrail_name=MOCK_REGISTER_REQUEST.guardrail_name, + status="pending_review", + submitted_at=datetime.now(), + ) + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + user = UserAPIKeyAuth(user_id="u1", user_email="alice@co.com", team_id="team-1") + + await register_guardrail(MOCK_REGISTER_REQUEST, user) + + stored_text = mock_prisma.db.litellm_guardrailstable.create.call_args[1]["data"]["litellm_params"] + assert "guardrails.example.com" not in stored_text + stored = json.loads(stored_text) + assert decrypt_if_encrypted_with(stored["api_base"], endpoint_salt_key) == "https://guardrails.example.com/validate" + assert decrypt_if_encrypted_with(stored["guardrail"], endpoint_salt_key) == "generic_guardrail_api" + + +@pytest.mark.asyncio +async def test_approve_guardrail_submission_initializes_the_decrypted_params(mocker, endpoint_salt_key): + from litellm.proxy.guardrails.guardrail_registry import encrypted_guardrail_litellm_params + + plaintext_params = { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "api_base": "https://g.com", + "api_key": "vendor-secret", + } + stored_params = json.loads(encrypted_guardrail_litellm_params(plaintext_params)) + assert stored_params["api_key"] != "vendor-secret" + mock_prisma = mocker.Mock() + row = mocker.Mock( + guardrail_id="approve-encrypted", + guardrail_name="my-guard", + status="pending_review", + litellm_params=stored_params, + guardrail_info={}, + team_id="team-1", + ) + mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row) + mock_prisma.db.litellm_guardrailstable.update = AsyncMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + mock_handler = mocker.Mock() + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_handler) + + await approve_guardrail_submission("approve-encrypted", UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)) + + initialized = mock_handler.initialize_guardrail.call_args.kwargs["guardrail"] + assert initialized["litellm_params"] == plaintext_params + + +@pytest.mark.asyncio +async def test_list_guardrails_v2_decrypts_a_stored_row_before_masking(mocker, endpoint_salt_key): + from litellm.proxy.guardrails.guardrail_registry import encrypted_guardrail_litellm_params + + stored_params = json.loads( + encrypted_guardrail_litellm_params( + { + "guardrail": "azure/text_moderations", + "mode": "pre_call", + "api_key": "sk-1234567890abcdef", + "api_base": "https://api.secret.example.com", + } + ) + ) + assert stored_params["guardrail"] != "azure/text_moderations" + db_row = { + "guardrail_id": "encrypted-db-guardrail", + "guardrail_name": "Encrypted DB Guardrail", + "litellm_params": stored_params, + "guardrail_info": {"description": "Test guardrail"}, + "created_at": datetime.now(), + "updated_at": datetime.now(), + } + mock_prisma_client = mocker.Mock() + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[db_row]) + mock_in_memory_handler = mocker.Mock() + mock_in_memory_handler.list_in_memory_guardrails.return_value = [] + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + + response = await list_guardrails_v2(user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)) + + assert len(response.guardrails) == 1 + litellm_params = response.guardrails[0].litellm_params + params = litellm_params if isinstance(litellm_params, dict) else litellm_params.model_dump() + assert params["api_key"] != "sk-1234567890abcdef" + assert "****" in str(params["api_key"]) + assert params["guardrail"] == "azure/text_moderations" + assert params["mode"] == "pre_call" + assert params["api_base"] == "https://api.secret.example.com" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 836668de0c8..e379ea8212d 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -1,11 +1,16 @@ -from collections.abc import Iterable +import json +from collections.abc import Iterable, Iterator from unittest.mock import AsyncMock, MagicMock import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy import proxy_server +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with from litellm.proxy.guardrails.guardrail_registry import ( + encrypted_guardrail_litellm_params, get_guardrail_initializer_from_hooks, + guardrail_from_row, GuardrailRegistry, InMemoryGuardrailHandler, ) @@ -1084,3 +1089,114 @@ def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance(): finally: for cb_list, snapshot in zip(lists, snapshots): cb_list[:] = snapshot + + +GUARDRAIL_SALT_KEY = "sk-guardrail-salt-1234" + + +class _StoredGuardrailRow: + def __init__(self, **fields: object) -> None: + self._fields = fields + + def __getattr__(self, name: str) -> object: + try: + return self._fields[name] + except KeyError as e: + raise AttributeError(name) from e + + def __iter__(self) -> Iterator[tuple[str, object]]: + return iter(self._fields.items()) + + +@pytest.fixture +def guardrail_salt_key(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", GUARDRAIL_SALT_KEY) + monkeypatch.setattr(proxy_server, "general_settings", {}) + return GUARDRAIL_SALT_KEY + + +def _vendor_params() -> LitellmParams: + return LitellmParams( + guardrail="openai_moderation", mode="pre_call", api_key="vendor-secret", api_base="https://vendor.example" + ) + + +@pytest.mark.asyncio +async def test_add_guardrail_to_db_stores_every_litellm_params_string_as_ciphertext(guardrail_salt_key): + prisma_client = MagicMock() + prisma_client.db.litellm_guardrailstable.create = AsyncMock(return_value=_StoredGuardrailRow(guardrail_id="g-1")) + guardrail = Guardrail(guardrail_name="moderation", litellm_params=_vendor_params()) + + created = await GuardrailRegistry().add_guardrail_to_db(guardrail=guardrail, prisma_client=prisma_client) + + stored_text = prisma_client.db.litellm_guardrailstable.create.call_args.kwargs["data"]["litellm_params"] + assert "vendor-secret" not in stored_text + assert "vendor.example" not in stored_text + stored = json.loads(stored_text) + assert decrypt_if_encrypted_with(stored["api_key"], guardrail_salt_key) == "vendor-secret" + assert decrypt_if_encrypted_with(stored["guardrail"], guardrail_salt_key) == "openai_moderation" + assert stored["default_on"] is False + assert created["guardrail_id"] == "g-1" + + +@pytest.mark.asyncio +async def test_update_guardrail_in_db_stores_ciphertext_and_returns_the_decrypted_guardrail(guardrail_salt_key): + prisma_client = MagicMock() + + async def _update(where: dict, data: dict) -> _StoredGuardrailRow: + return _StoredGuardrailRow( + guardrail_id=where["guardrail_id"], + guardrail_name=data["guardrail_name"], + litellm_params=json.loads(data["litellm_params"]), + guardrail_info={}, + status="active", + ) + + prisma_client.db.litellm_guardrailstable.update = AsyncMock(side_effect=_update) + guardrail = Guardrail(guardrail_name="moderation", litellm_params=_vendor_params()) + + updated = await GuardrailRegistry().update_guardrail_in_db( + guardrail_id="g-1", guardrail=guardrail, prisma_client=prisma_client + ) + + stored = json.loads(prisma_client.db.litellm_guardrailstable.update.call_args.kwargs["data"]["litellm_params"]) + assert decrypt_if_encrypted_with(stored["api_key"], guardrail_salt_key) == "vendor-secret" + assert updated["litellm_params"]["api_key"] == "vendor-secret" + assert updated["litellm_params"]["guardrail"] == "openai_moderation" + + +@pytest.mark.asyncio +async def test_get_guardrail_by_id_from_db_decrypts_encrypted_rows_and_reads_legacy_plaintext_rows( + guardrail_salt_key, +): + plaintext_params = {"guardrail": "openai_moderation", "mode": "pre_call", "api_key": "vendor-secret"} + encrypted_row = _StoredGuardrailRow( + guardrail_id="g-encrypted", + guardrail_name="encrypted", + litellm_params=json.loads(encrypted_guardrail_litellm_params(plaintext_params)), + guardrail_info={}, + status="active", + ) + legacy_row = _StoredGuardrailRow( + guardrail_id="g-legacy", + guardrail_name="legacy", + litellm_params={**plaintext_params, "api_key": "legacy-secret"}, + guardrail_info={}, + status="active", + ) + assert encrypted_row.litellm_params["api_key"] != "vendor-secret" + prisma_client = MagicMock() + prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(side_effect=[encrypted_row, legacy_row]) + registry = GuardrailRegistry() + + encrypted = await registry.get_guardrail_by_id_from_db(guardrail_id="g-encrypted", prisma_client=prisma_client) + legacy = await registry.get_guardrail_by_id_from_db(guardrail_id="g-legacy", prisma_client=prisma_client) + + assert encrypted is not None and encrypted["litellm_params"] == plaintext_params + assert legacy is not None and legacy["litellm_params"]["api_key"] == "legacy-secret" + + +def test_guardrail_from_row_keeps_a_row_without_litellm_params(guardrail_salt_key): + row = _StoredGuardrailRow(guardrail_id="g-empty", guardrail_name="empty", litellm_params=None, guardrail_info=None) + + assert guardrail_from_row(row)["litellm_params"] is None diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index db87e12ac88..b5588dd6e5e 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -184,6 +184,28 @@ async def test_detail_db_row_still_resolves(): assert resp.type == "ContentSafety" +@pytest.mark.asyncio +async def test_detail_encrypted_db_row_resolves_its_provider(monkeypatch): + import json + + from litellm.proxy import proxy_server + from litellm.proxy.guardrails.guardrail_registry import encrypted_guardrail_litellm_params + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-usage-salt-1234") + monkeypatch.setattr(proxy_server, "general_settings", {}) + row = _db_row(guardrail_id="db-encrypted", provider="aim") + row.litellm_params = json.loads(encrypted_guardrail_litellm_params({"guardrail": "aim", "mode": "pre_call"})) + assert row.litellm_params["guardrail"] != "aim" + prisma = _prisma(find_unique=row) + handler = _config_handler() + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="db-encrypted", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert resp.provider == "aim" + + # ---- overview --------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py index 0ecc4f8d7cb..5fa096b1523 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py +++ b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py @@ -458,6 +458,27 @@ async def test_scan_covered_tables_classifies_legacy_and_v2(salt_key, monkeypatc assert by_loc["credentials"].legacy == 0 +@pytest.mark.asyncio +async def test_scan_covered_tables_reads_nested_guardrail_params(salt_key, monkeypatch): + legacy = _legacy_ct("guardrail-secret", monkeypatch) + + client = MagicMock() + _empty_covered_tables(client) + client.db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[ + SimpleNamespace( + litellm_params={"guardrail": "openai_moderation", "api_key": legacy, "extra": {"nested": legacy}} + ) + ] + ) + client.db.litellm_config.find_unique = AsyncMock(return_value=None) + + by_loc = {r.location: r for r in await cm._scan_covered_tables(client)} + + assert by_loc["guardrails"].legacy == 2 + assert by_loc["guardrails"].plaintext == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("column", ("static_headers", "env")) @pytest.mark.parametrize("algorithm", ("xsalsa20-poly1305", "aes-256-gcm")) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index aa6be328f4a..f460141f581 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -21093,3 +21093,66 @@ class TestTeamAdminMemberKeyBudgetUpdate: ) assert exc.value.status_code == 403 assert "member_key_budgets" not in str(exc.value.detail) + + +@pytest.mark.asyncio +@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key") +async def test_rotate_master_key_rekeys_router_settings_and_guardrail_params(mock_rotate_mcp, monkeypatch): + import prisma + + from litellm.proxy import proxy_server + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_json_strings + from litellm.proxy.management_endpoints.key_management_endpoints import _rotate_master_key + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key") + monkeypatch.setattr(proxy_server, "general_settings", {}) + stored_router_settings = encrypt_json_strings({"redis_password": "redis-pw", "num_retries": 2}) + stored_guardrail_params = encrypt_json_strings( + {"guardrail": "openai_moderation", "api_key": "vendor-secret", "default_on": False} + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.tx = MagicMock( + return_value=AsyncMock(__aenter__=AsyncMock(return_value=AsyncMock()), __aexit__=AsyncMock(return_value=False)) + ) + mock_prisma_client.db.litellm_config.find_many = AsyncMock( + return_value=[SimpleNamespace(param_name="router_settings", param_value=stored_router_settings)] + ) + mock_prisma_client.db.litellm_config.update = AsyncMock() + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[SimpleNamespace(guardrail_id="guardrail-1", litellm_params=stored_guardrail_params)] + ) + mock_prisma_client.db.litellm_guardrailstable.update = AsyncMock() + mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + mock_rotate_mcp.return_value = None + mock_proxy_config = MagicMock() + mock_proxy_config.decrypt_model_list_from_db.return_value = [] + + with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config): + await _rotate_master_key( + prisma_client=mock_prisma_client, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="test-user" + ), + current_master_key="sk-old-master-key", + new_master_key="sk-new-master-key", + ) + + router_update = mock_prisma_client.db.litellm_config.update.call_args + assert router_update.kwargs["where"] == {"param_name": "router_settings"} + assert isinstance(router_update.kwargs["data"]["param_value"], prisma.Json) + rotated_router_settings = router_update.kwargs["data"]["param_value"].data + assert decrypt_if_encrypted_with(rotated_router_settings["redis_password"], "sk-new-master-key") == "redis-pw" + assert decrypt_if_encrypted_with(rotated_router_settings["redis_password"], "sk-old-master-key") is None + assert rotated_router_settings["num_retries"] == 2 + + guardrail_update = mock_prisma_client.db.litellm_guardrailstable.update.call_args + assert guardrail_update.kwargs["where"] == {"guardrail_id": "guardrail-1"} + assert isinstance(guardrail_update.kwargs["data"]["litellm_params"], prisma.Json) + rotated_guardrail_params = guardrail_update.kwargs["data"]["litellm_params"].data + assert decrypt_if_encrypted_with(rotated_guardrail_params["api_key"], "sk-new-master-key") == "vendor-secret" + assert decrypt_if_encrypted_with(rotated_guardrail_params["guardrail"], "sk-new-master-key") == "openai_moderation" + assert rotated_guardrail_params["default_on"] is False diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 5f7807650e1..289d1438467 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -23,7 +23,7 @@ from litellm.proxy._types import ( ReconcileOutcome, UserAPIKeyAuth, ) -from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, _get_team_deployments, @@ -1397,8 +1397,8 @@ class TestUpdateModel: new=AsyncMock(return_value=None), ), patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", - side_effect=lambda value: value, + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, ), patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -1450,8 +1450,8 @@ class TestUpdateModel: new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", - side_effect=lambda value: value, + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, ), patch( "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -2547,7 +2547,7 @@ class TestAddAndDeleteModelLifecycle: mock_router.get_model_ids.return_value = [model_id] _PS = "litellm.proxy.proxy_server" - _ENCRYPT = "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper" + _ENCRYPT = "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper" with ( patch(f"{_PS}.prisma_client", mock_prisma), patch(f"{_PS}.store_model_in_db", True), @@ -4145,7 +4145,7 @@ class TestModelInfoServerDerivedPricingFilter: mock_router.get_model_ids.return_value = [model_id] _PS = "litellm.proxy.proxy_server" - _ENCRYPT = "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper" + _ENCRYPT = "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper" with ( patch(f"{_PS}.prisma_client", mock_prisma), patch(f"{_PS}.store_model_in_db", True), @@ -4530,7 +4530,7 @@ class TestUpdateDBModelClearCredentialName: litellm_params=updateLiteLLMParams(litellm_credential_name=None) ) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4559,7 +4559,7 @@ class TestUpdateDBModelClearCredentialName: ) update_patch: Final = updateDeployment(litellm_params=updateLiteLLMParams(tpm=10)) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4578,7 +4578,7 @@ class TestUpdateDBModelClearCredentialName: litellm_params=updateLiteLLMParams(litellm_credential_name=None) ) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4606,7 +4606,7 @@ class TestUpdateDBModelClearCredentialName: ) ) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4633,7 +4633,7 @@ class TestUpdateDBModelClearCredentialName: litellm_params=updateLiteLLMParams(litellm_credential_name="other-credential") ) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4709,7 +4709,7 @@ class TestPatchModelCredentialName: new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), patch( - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value, ), patch( @@ -6358,7 +6358,7 @@ class TestStrategyRouterWriteValidation: new=AsyncMock(return_value=None), ), patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", lambda value, new_encryption_key=None: value, ), ): @@ -6424,7 +6424,7 @@ class TestStrategyRouterWriteValidation: ) with ( patch( # test-quality-ok: params are encrypted with the proxy master key, which this test does not configure - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", lambda value, new_encryption_key=None: value, ), patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted @@ -6464,7 +6464,7 @@ class TestStrategyRouterWriteValidation: new=AsyncMock(return_value=None), ), patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", lambda value, new_encryption_key=None: value, ), ): @@ -7221,7 +7221,7 @@ class TestAccessGroupModelSync: f"{self._MOD}.clear_cache", new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), - patch(f"{self._MOD}.encrypt_value_helper", side_effect=lambda value, **kwargs: value), + patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value), ): stack.enter_context(target) yield stack.enter_context(patch(self._INVALIDATE, new=AsyncMock())) @@ -7544,6 +7544,14 @@ class TestTeamMemberAutoRouterWrites: assert saved_info["team_id"] == "member-team" assert saved_info["access_groups"] == ["retained-admin-group"] + @staticmethod + def _with_decrypted_jev_key(config: Mapping[str, object]) -> dict[str, object]: + jev: Final = config.get("jev_classifier_config") + if not isinstance(jev, dict) or not isinstance(jev.get("api_key"), str): + return dict(config) + decrypted_key: Final = decrypt_value_helper(jev["api_key"], key="api_key") + return {**config, "jev_classifier_config": {**jev, "api_key": decrypted_key}} + @pytest.mark.asyncio @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) @pytest.mark.parametrize("change", ["save", "rotate", "move", "move-without-key", "reset", "heuristic"]) @@ -7599,7 +7607,10 @@ class TestTeamMemberAutoRouterWrites: if change == "heuristic" else {**config, "jev_classifier_config": {**transport, "timeout_ms": 8100, **overrides}} ) - assert saved == expected + stored_jev: Final = saved.get("jev_classifier_config") + if isinstance(stored_jev, dict) and isinstance(stored_jev.get("api_key"), str): + assert stored_jev["api_key"] != expected["jev_classifier_config"]["api_key"] + assert self._with_decrypted_jev_key(saved) == expected assert row.litellm_params["complexity_router_config"] == stored_config assert request.litellm_params.complexity_router_config == config @@ -7824,8 +7835,8 @@ class TestModelManagementActorEdges: patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", - side_effect=lambda value: value, + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, ), patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -7873,8 +7884,8 @@ class TestModelManagementActorEdges: patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", - side_effect=lambda value: value, + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, ), patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -7994,8 +8005,8 @@ class TestModelManagementActorEdges: patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation - "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", - side_effect=lambda value: value, + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, ), patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -8051,3 +8062,122 @@ class TestModelManagementActorEdges: assert response.status_code == 400 assert "Cannot edit config-based model" in response.text prisma.db.litellm_proxymodeltable.update.assert_not_awaited() + + +class TestNestedLitellmParamsEncryption: + SALT_KEY: Final = "sk-nested-salt-1234" + + @pytest.fixture(autouse=True) + def _salt_key(self, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", self.SALT_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + @pytest.mark.asyncio + async def test_add_new_model_encrypts_nested_extra_headers(self): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + model_id: Final = "nested-headers-model" + db_row: Final = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="gateway-model", + litellm_params={"model": "openai/gpt-5.4-mini"}, + model_info={"id": model_id}, + created_by="test-admin", + updated_by="test-admin", + ) + mock_prisma: Final = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=db_row) + mock_proxy_config: Final = MagicMock() + mock_proxy_config._add_deployment_locked = AsyncMock( + return_value=ReconcileOutcome(still_desired=frozenset(), live_after=frozenset()) + ) + mock_router: Final = MagicMock() + mock_router.get_model_ids.return_value = [model_id] + _PS: Final = "litellm.proxy.proxy_server" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.proxy_config", mock_proxy_config), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", mock_router), + ): + await add_new_model( + model_params=Deployment( + model_name="gateway-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-5.4-mini", + api_key="sk-placeholder", + extra_headers={"Authorization": "Bearer gateway-secret", "X-Gateway-Token": "gw-token"}, + rpm=10, + ), + model_info={"id": model_id}, + ), + user_api_key_dict=UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + create_call: Final = mock_prisma.db.litellm_proxymodeltable.create.call_args + stored_text: Final = (create_call.kwargs["data"] if "data" in create_call.kwargs else create_call.args[0])[ + "litellm_params" + ] + assert "gateway-secret" not in stored_text + assert "gw-token" not in stored_text + stored: Final = json.loads(stored_text) + assert decrypt_if_encrypted_with(stored["extra_headers"]["Authorization"], self.SALT_KEY) == "Bearer gateway-secret" + assert decrypt_if_encrypted_with(stored["extra_headers"]["X-Gateway-Token"], self.SALT_KEY) == "gw-token" + assert decrypt_if_encrypted_with(stored["api_key"], self.SALT_KEY) == "sk-placeholder" + assert stored["rpm"] == 10 + + def test_update_db_model_encrypts_nested_extra_headers_and_reencrypts_a_legacy_plaintext_row(self): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gateway-model", + litellm_params=LiteLLM_Params(model="openai/gpt-5.4-mini", api_key="legacy-plaintext-key"), + model_info={"id": "legacy-model"}, + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(extra_headers={"Authorization": "Bearer gateway-secret"}) + ) + + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + assert "gateway-secret" not in result["litellm_params"] + assert "legacy-plaintext-key" not in result["litellm_params"] + stored: Final = json.loads(result["litellm_params"]) + assert decrypt_if_encrypted_with(stored["extra_headers"]["Authorization"], self.SALT_KEY) == "Bearer gateway-secret" + assert decrypt_if_encrypted_with(stored["api_key"], self.SALT_KEY) == "legacy-plaintext-key" + assert decrypt_if_encrypted_with(stored["model"], self.SALT_KEY) == "openai/gpt-5.4-mini" + + def test_update_db_model_keeps_an_already_encrypted_value_decryptable_after_a_second_write(self): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + first: Final = update_db_model( + db_model=Deployment( + model_name="gateway-model", + litellm_params=LiteLLM_Params(model="openai/gpt-5.4-mini", api_key="sk-first"), + model_info={"id": "twice-written"}, + ), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(extra_headers={"Authorization": "Bearer gateway-secret"}) + ), + ) + second: Final = update_db_model( + db_model=Deployment( + model_name="gateway-model", + litellm_params=LiteLLM_Params(**json.loads(first["litellm_params"])), + model_info={"id": "twice-written"}, + ), + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(rpm=5)), + ) + + stored: Final = json.loads(second["litellm_params"]) + assert decrypt_if_encrypted_with(stored["extra_headers"]["Authorization"], self.SALT_KEY) == "Bearer gateway-secret" + assert decrypt_if_encrypted_with(stored["api_key"], self.SALT_KEY) == "sk-first" + assert stored["rpm"] == 5 diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index b2ef327f50e..0726cba760e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -1111,6 +1111,25 @@ async def test_ProxyConfig_save_config_writes_only_changed_router_settings(monke assert table.upserted_param_names == ["router_settings"] +@pytest.mark.asyncio +async def test_ProxyConfig_save_config_encrypts_router_settings_and_reencrypts_the_legacy_row(monkeypatch): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-save-config-salt-1234") + proxy_config, table = _db_backed_proxy_config(monkeypatch, {"router_settings": {"redis_host": "redis.internal"}}) + baseline: Final = {"model_list": [], "router_settings": {"num_retries": 1}} + proxy_config.update_config_state(config=baseline) + changed: Final = {**baseline, "router_settings": {"num_retries": 1, "redis_password": "redis-canary-password"}} + + await proxy_config.save_config(changed) + + stored: Final = table.rows["router_settings"] + assert "redis-canary-password" not in json.dumps(stored) + assert "redis.internal" not in json.dumps(stored) + assert decrypt_json_strings(stored) == {"redis_host": "redis.internal", "redis_password": "redis-canary-password"} + assert table.upserted_param_names == ["router_settings"] + + @pytest.mark.asyncio async def test_ProxyConfig_save_config_removes_a_key_only_when_the_db_has_it(monkeypatch): proxy_config, table = _db_backed_proxy_config( diff --git a/tests/test_litellm/proxy/test_fallback_management_endpoints.py b/tests/test_litellm/proxy/test_fallback_management_endpoints.py index 054dafbf5a7..d701a3e0d88 100644 --- a/tests/test_litellm/proxy/test_fallback_management_endpoints.py +++ b/tests/test_litellm/proxy/test_fallback_management_endpoints.py @@ -346,6 +346,39 @@ class TestCreateFallback: # Verify the correct attribute was updated assert hasattr(mock_router, "context_window_fallbacks") + async def test_create_fallback_encrypts_the_router_settings_row( + self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict, monkeypatch + ): + import json + + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-fallback-salt-1234") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + mock_proxy_config.get_config = AsyncMock( + return_value={"router_settings": {"redis_password": "redis-canary-password"}} + ) + request = FallbackCreateRequest( + model="gpt-3.5-turbo", fallback_models=["gpt-4", "claude-3-haiku"], fallback_type="general" + ) + + with ( + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + await create_fallback(request, mock_user_api_key_dict) + + upsert_data = mock_prisma_client.db.litellm_config.upsert.call_args.kwargs["data"] + stored_text = upsert_data["create"]["param_value"] + assert stored_text == upsert_data["update"]["param_value"] + assert "redis-canary-password" not in stored_text + assert decrypt_json_strings(json.loads(stored_text)) == { + "redis_password": "redis-canary-password", + "fallbacks": [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}], + } + @pytest.mark.asyncio class TestGetFallback: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 884a9c81500..fe5497dcb31 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -15368,3 +15368,56 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is call(api_provider="openai", capture_rate=0.97), call(api_provider="openai", capture_rate=None), ] + + +ROUTER_SETTINGS_SALT_KEY = "sk-router-settings-salt-1234" + + +def test_update_config_router_settings_encrypted_before_write(_update_config_setup, monkeypatch): + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_if_encrypted_with, + encrypt_json_strings, + ) + + monkeypatch.setenv("LITELLM_SALT_KEY", ROUTER_SETTINGS_SALT_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + client, prisma, restore = _update_config_setup( + initial_rows={"router_settings": encrypt_json_strings({"redis_host": "redis.internal"})} + ) + try: + resp = client.post("/config/update", json={"router_settings": {"redis_password": "redis-canary-password"}}) + assert resp.status_code == 200, resp.text + stored = prisma.db.litellm_config.rows["router_settings"] + assert "redis-canary-password" not in json.dumps(stored) + assert "redis.internal" not in json.dumps(stored) + assert decrypt_if_encrypted_with(stored["redis_password"], ROUTER_SETTINGS_SALT_KEY) == "redis-canary-password" + assert decrypt_if_encrypted_with(stored["redis_host"], ROUTER_SETTINGS_SALT_KEY) == "redis.internal" + finally: + restore() + + +@pytest.mark.asyncio +async def test_add_router_settings_from_db_config_decrypts_the_stored_row(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_json_strings + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("LITELLM_SALT_KEY", ROUTER_SETTINGS_SALT_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + proxy_config = ProxyConfig() + mock_router = MagicMock() + mock_router.update_settings = MagicMock() + stored_row = MagicMock() + stored_row.param_value = encrypt_json_strings({"redis_password": "redis-canary-password", "num_retries": 3}) + assert stored_row.param_value["redis_password"] != "redis-canary-password" + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=stored_row) + proxy_config.router_settings.load_yaml({"routing_strategy": "simple-shuffle"}) + + await proxy_config._add_router_settings_from_db_config(llm_router=mock_router, prisma_client=mock_prisma_client) + + combined_settings = mock_router.update_settings.call_args[1] + assert combined_settings["redis_password"] == "redis-canary-password" + assert combined_settings["num_retries"] == 3 + assert combined_settings["routing_strategy"] == "simple-shuffle" diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index e185d95ffb8..306aa95a4a5 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -328,14 +328,11 @@ class TestModelRepository: @pytest.mark.asyncio @patch( - "litellm.repositories.model_repository.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: f"encrypted_{v}", ) - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_create_model_encrypts_params(self, mock_decrypt, mock_encrypt, repo): + async def test_create_model_encrypts_params(self, mock_encrypt): + repo = ModelRepository(MockPrismaClient(), encryption_key="unit-test-key") model = await repo.create_model( model_name="gpt-4", litellm_params={"api_key": "sk-secret"}, @@ -346,14 +343,10 @@ class TestModelRepository: @pytest.mark.asyncio @patch( - "litellm.repositories.model_repository.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: f"encrypted_{v}", ) - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_create_model_all_fields(self, mock_decrypt, mock_encrypt, repo): + async def test_create_model_all_fields(self, mock_encrypt, repo): model = await repo.create_model( model_name="gpt-4-turbo", litellm_params={ @@ -370,14 +363,10 @@ class TestModelRepository: @pytest.mark.asyncio @patch( - "litellm.repositories.model_repository.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: f"encrypted_{v}", ) - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_update_model_all_fields(self, mock_decrypt, mock_encrypt, repo): + async def test_update_model_all_fields(self, mock_encrypt, repo): repo._prisma_client.db.litellm_proxymodeltable._records["model-full"] = { "model_id": "model-full", "model_name": "old-name", @@ -395,11 +384,7 @@ class TestModelRepository: assert updated.model_name == "new-name" @pytest.mark.asyncio - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_find_all(self, mock_decrypt, repo): + async def test_find_all(self, repo): repo._prisma_client.db.litellm_proxymodeltable._records = { "m1": { "model_id": "m1", @@ -418,11 +403,7 @@ class TestModelRepository: assert len(models) == 2 @pytest.mark.asyncio - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_find_unblocked(self, mock_decrypt, repo): + async def test_find_unblocked(self, repo): repo._prisma_client.db.litellm_proxymodeltable._records = { "m1": { "model_id": "m1", @@ -435,11 +416,7 @@ class TestModelRepository: assert len(models) == 1 @pytest.mark.asyncio - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_find_by_name(self, mock_decrypt, repo): + async def test_find_by_name(self, repo): repo._prisma_client.db.litellm_proxymodeltable._records = { "m1": { "model_id": "m1", @@ -452,14 +429,10 @@ class TestModelRepository: @pytest.mark.asyncio @patch( - "litellm.repositories.model_repository.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: v, ) - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_update_model(self, mock_decrypt, mock_encrypt, repo): + async def test_update_model(self, mock_encrypt, repo): repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { "model_id": "m1", "model_name": "gpt-4", @@ -474,11 +447,7 @@ class TestModelRepository: assert updated.blocked is True @pytest.mark.asyncio - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_delete_model(self, mock_decrypt, repo): + async def test_delete_model(self, repo): repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { "model_id": "m1", "model_name": "gpt-4", @@ -489,14 +458,10 @@ class TestModelRepository: @pytest.mark.asyncio @patch( - "litellm.repositories.model_repository.encrypt_value_helper", + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: v, ) - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda v, **kw: v, - ) - async def test_block_unblock_model(self, mock_decrypt, mock_encrypt, repo): + async def test_block_unblock_model(self, mock_encrypt, repo): repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { "model_id": "m1", "model_name": "gpt-4", @@ -1894,11 +1859,7 @@ class TestModelRepositoryExtended: return ModelRepository(client) @pytest.mark.asyncio - @patch( - "litellm.repositories.model_repository.decrypt_value_helper", - side_effect=lambda value, **kw: value, - ) - async def test_find_by_team_id(self, mock_decrypt, repo): + async def test_find_by_team_id(self, repo): repo._prisma_client.db.litellm_proxymodeltable._records["model-1"] = { "model_id": "model-1", "model_name": "gpt-4", @@ -2290,3 +2251,42 @@ class TestAutoRouterSessionRepository: assert AutoRouterSessionRepository(client).table is client.db.litellm_autoroutersession with pytest.raises(RuntimeError, match="No DB Connected"): _ = AutoRouterSessionRepository(None).table + + +class TestModelRepositoryNestedEncryption: + SALT_KEY: Final = "sk-repository-salt-1234" + + @pytest.fixture(autouse=True) + def _salt_key(self, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", self.SALT_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + @pytest.mark.asyncio + async def test_create_model_encrypts_nested_strings_and_find_by_id_decrypts_them(self): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with + + client: Final = MockPrismaClient() + repo: Final = ModelRepository(client, encryption_key=self.SALT_KEY) + litellm_params: Final = { + "model": "openai/gpt-5.4-mini", + "api_key": "sk-placeholder", + "extra_headers": {"Authorization": "Bearer gateway-secret", "X-Gateway-Token": "gw-token"}, + "rpm": 10, + } + + created: Final = await repo.create_model( + model_name="gateway-model", litellm_params=litellm_params, created_by="admin", model_id="gw-1" + ) + + stored_text: Final = client.db.litellm_proxymodeltable._records["gw-1"]["litellm_params"] + assert "gateway-secret" not in stored_text + assert "gw-token" not in stored_text + stored: Final = json.loads(stored_text) + assert decrypt_if_encrypted_with(stored["extra_headers"]["Authorization"], self.SALT_KEY) == "Bearer gateway-secret" + assert decrypt_if_encrypted_with(stored["api_key"], self.SALT_KEY) == "sk-placeholder" + assert stored["rpm"] == 10 + assert created is not None + assert created.litellm_params == litellm_params + found: Final = await repo.find_by_id("gw-1") + assert found is not None + assert found.litellm_params == litellm_params From 35f3b16c950b1b8732ba3fd2875084045e8103f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:38:41 -0700 Subject: [PATCH 2/7] fix(proxy): re-key string-typed guardrail rows on rotation, skip malformed rows, keep secrets out of debug logs Master key rotation now parses a guardrail or config row stored as a JSON string before re-keying it and skips a row that is not a JSON object with an error naming the row instead of raising or rewriting it, and _guardrail_rows only treats a missing table as empty so a lost connection surfaces. A stored value that is not a JSON object logs one error without its contents when read. ModelRepository reads decrypt with the same injected key its writes use. The guardrail registry no longer logs the full litellm_params document at debug and reports only the names of changed fields. lite encryption migrate --check scans the router_settings config row next to environment_variables. --- .../common_utils/encrypt_decrypt_utils.py | 36 +++-- .../proxy/guardrails/guardrail_registry.py | 3 +- .../credential_migration.py | 31 +++-- .../key_management_endpoints.py | 61 ++++++--- litellm/repositories/model_repository.py | 5 +- .../test_encrypt_decrypt_utils.py | 38 ++++++ .../guardrails/test_guardrail_registry.py | 45 +++++++ .../test_credential_migration.py | 19 +++ .../test_key_management_endpoints.py | 125 +++++++++++++++--- tests/unit/repositories/test_repositories.py | 15 +++ 10 files changed, 313 insertions(+), 65 deletions(-) diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index b83aa287ba9..7b8a9fcb383 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -274,11 +274,11 @@ def json_value(value: object) -> JsonValue: return _JSON_VALUE.validate_json(safe_dumps(plain)) -def _plaintext_of(value: str) -> str: - signing_key: Final = _get_salt_key() - if signing_key is None: +def _plaintext_of(value: str, signing_key: str | None = None) -> str: + key: Final = signing_key or _get_salt_key() + if key is None: return value - decrypted: Final = decrypt_if_encrypted_with(value, signing_key) + decrypted: Final = decrypt_if_encrypted_with(value, key) return value if decrypted is None else decrypted @@ -304,25 +304,41 @@ def encrypt_json_strings(value: JsonValue, new_encryption_key: str | None = None return value -def decrypt_json_strings(value: JsonValue, depth: int = 0) -> JsonValue: +def decrypt_json_strings(value: JsonValue, signing_key: str | None = None, depth: int = 0) -> JsonValue: if depth > DEFAULT_MAX_RECURSE_DEPTH: return value match value: case str(): - return _plaintext_of(value) + return _plaintext_of(value, signing_key) case list(): - return [decrypt_json_strings(item, depth + 1) for item in value] + return [decrypt_json_strings(item, signing_key, depth + 1) for item in value] case dict(): - return {key: decrypt_json_strings(item, depth + 1) for key, item in value.items()} + return {key: decrypt_json_strings(item, signing_key, depth + 1) for key, item in value.items()} case _: return value -def _stored_json_object(value: object) -> dict[str, JsonValue]: +def parse_stored_json_object(value: object) -> dict[str, JsonValue] | None: try: return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value) except ValidationError: - return {} + return None + + +def _stored_json_object(value: object) -> dict[str, JsonValue]: + parsed: Final = parse_stored_json_object(value) + if parsed is not None: + return parsed + if value is not None: + verbose_proxy_logger.error("Stored value is not a JSON object (%s); reading it as empty", type(value).__name__) + return {} + + +def encrypt_stored_json_object(value: object, new_encryption_key: str | None = None) -> dict[str, JsonValue] | None: + parsed: Final = parse_stored_json_object(value) + if parsed is None: + return None + return {key: encrypt_json_strings(item, new_encryption_key=new_encryption_key) for key, item in parsed.items()} def decrypt_stored_json_object(value: object) -> dict[str, JsonValue]: diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 5ddc9238d1c..982b15a1ffa 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -531,7 +531,6 @@ class InMemoryGuardrailHandler: return self.IN_MEMORY_GUARDRAILS[guardrail_id] litellm_params_data: Final = guardrail["litellm_params"] - verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data) if isinstance(litellm_params_data, dict): litellm_params = LitellmParams(**litellm_params_data) @@ -820,7 +819,7 @@ class InMemoryGuardrailHandler: # Log differences if any found if changed_fields: - verbose_proxy_logger.debug("Guardrail params changed. Differences: %s", changed_fields) + verbose_proxy_logger.debug("Guardrail params changed: %s", sorted(changed_fields)) # Return True if any fields changed return len(changed_fields) > 0 diff --git a/litellm/proxy/management_endpoints/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py index 7ab8966549c..b6645574e31 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -48,6 +48,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decode_secret_map, decrypt_value_helper, encrypt_value_helper, + parse_stored_json_object, ) ValueClass = Literal["migrated", "legacy", "plaintext", "undecryptable", "not-a-string"] @@ -541,23 +542,23 @@ async def _scan_one_table( return report -async def _scan_config_env_vars(prisma_client: object) -> LocationReport: - """Scan the ``environment_variables`` config row (``param_value`` dict).""" - report: Final = LocationReport(location="config_environment_variables") +_CONFIG_SECTION_LOCATIONS: Final = ( + ("environment_variables", "config_environment_variables"), + ("router_settings", "config_router_settings"), +) + + +async def _scan_config_section(prisma_client: object, param_name: str, location: str) -> LocationReport: + """Scan one encrypted ``LiteLLM_Config`` row (``param_value`` JSON object).""" + report: Final = LocationReport(location=location) try: - record: Final = await prisma_client.db.litellm_config.find_unique(where={"param_name": "environment_variables"}) + record: Final = await prisma_client.db.litellm_config.find_unique(where={"param_name": param_name}) except Exception as e: # pragma: no cover - defensive - verbose_proxy_logger.debug("scan: config env vars unavailable: %s", str(e)) + verbose_proxy_logger.debug("scan: config %s unavailable: %s", param_name, str(e)) return report if record is None or record.param_value is None: return report - value = record.param_value - if isinstance(value, str): - try: - value = json.loads(value) - except (ValueError, TypeError): - value = {} - for s in _iter_encrypted_strings(value): + for s in _iter_encrypted_strings(parse_stored_json_object(record.param_value) or {}): _classify_into_report(report, s) return report @@ -567,7 +568,8 @@ async def _scan_covered_tables(prisma_client: object) -> list[LocationReport]: reports: Final[list[LocationReport]] = [] for location, db_attr, json_cols, scalar_cols in _COVERED_TABLE_SPECS: reports.append(await _scan_one_table(prisma_client, location, db_attr, json_cols, scalar_cols)) - reports.append(await _scan_config_env_vars(prisma_client)) + for param_name, location in _CONFIG_SECTION_LOCATIONS: + reports.append(await _scan_config_section(prisma_client, param_name, location)) return reports @@ -582,7 +584,8 @@ _CLOUDZERO_SENSITIVE: Final = ["api_key"] async def _migrate_covered_tables(prisma_client: object, user_api_key_dict: object) -> list[LocationReport]: """Re-encrypt the tables already covered by ``_rotate_master_key`` (model - table, credentials, MCP credential/env tables, config environment_variables) + table, credentials, guardrails, MCP credential/env tables, config + environment_variables and router_settings) by running that orchestrator in *same-key* mode. With the AES gate on, the re-encrypt writes land in ``v2:`` format. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5c4c369dd73..4132ddb708f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -83,8 +83,8 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( from litellm.proxy.common_utils.encrypt_decrypt_utils import ( ENCRYPTED_CONFIG_SECTIONS, encrypt_config_section, - encrypt_json_strings, - json_value, + encrypt_stored_json_object, + parse_stored_json_object, ) from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time @@ -5224,12 +5224,49 @@ async def delete_key_aliases( async def _guardrail_rows(prisma_client: PrismaClient) -> "Sequence[prisma_models.LiteLLM_GuardrailsTable]": + from prisma.errors import TableNotFoundError + try: return await GuardrailsRepository(prisma_client).table.find_many() - except Exception: # noqa: BLE001 # a deployment without the guardrails table still rotates every other store + except TableNotFoundError: return () +async def _rekey_config_section( + prisma_client: PrismaClient, param_name: str, param_value: object, new_master_key: str +) -> None: + import prisma + + section: Final = parse_stored_json_object(param_value) + if section is None: + verbose_proxy_logger.error( + "Master key rotation skipped config %s: stored value is not a JSON object", param_name + ) + return + await _config_table(prisma_client).update( + where={"param_name": param_name}, + data={ + "param_value": prisma.Json(encrypt_config_section(param_name, section, new_encryption_key=new_master_key)) + }, + ) + + +async def _rekey_guardrail_row( + prisma_client: PrismaClient, guardrail_row: "prisma_models.LiteLLM_GuardrailsTable", new_master_key: str +) -> None: + import prisma + + rekeyed: Final = encrypt_stored_json_object(guardrail_row.litellm_params, new_encryption_key=new_master_key) + if rekeyed is None: + verbose_proxy_logger.error( + "Master key rotation skipped guardrail %s: litellm_params is not a JSON object", guardrail_row.guardrail_id + ) + return + await GuardrailsRepository(prisma_client).table.update( + where={"guardrail_id": guardrail_row.guardrail_id}, data={"litellm_params": prisma.Json(rekeyed)} + ) + + async def _rotate_master_key( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, @@ -5317,27 +5354,13 @@ async def _rotate_master_key( for c in config: if c.param_name in ENCRYPTED_CONFIG_SECTIONS and c.param_value is not None: - await _config_table(prisma_client).update( - where={"param_name": c.param_name}, - data={ - "param_value": prisma.Json( - encrypt_config_section(c.param_name, c.param_value, new_encryption_key=new_master_key) - ) - }, - ) + await _rekey_config_section(prisma_client, c.param_name, c.param_value, new_master_key) # 3b. process guardrails table for guardrail_row in await _guardrail_rows(prisma_client): if guardrail_row.litellm_params is None: continue - await GuardrailsRepository(prisma_client).table.update( - where={"guardrail_id": guardrail_row.guardrail_id}, - data={ - "litellm_params": prisma.Json( - encrypt_json_strings(json_value(guardrail_row.litellm_params), new_encryption_key=new_master_key) - ) - }, - ) + await _rekey_guardrail_row(prisma_client, guardrail_row, new_master_key) # 4. process MCP server table try: diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index c21c8c4f653..1371a2b1a25 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -46,7 +46,10 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): } def _decrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]: - return {key: decrypt_json_strings(json_value(value)) for key, value in litellm_params.items()} + return { + key: decrypt_json_strings(json_value(value), signing_key=self._encryption_key) + for key, value in litellm_params.items() + } def _to_model(self, record: Any) -> LiteLLM_ProxyModelTable | None: """Convert a database record to a Model with decryption.""" diff --git a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py index d32f1c37a8b..6856ff46a44 100644 --- a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -23,6 +23,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_config_section, encrypt_json_strings, + encrypt_stored_json_object, encrypt_value, encrypt_value_helper, json_value, @@ -362,3 +363,40 @@ def test_encrypt_config_section_rekeys_router_settings_for_master_key_rotation() rotated = encrypt_config_section("router_settings", stored, new_encryption_key="sk-rotated") assert decrypt_if_encrypted_with(rotated["redis_password"], "sk-rotated") == "redis-pw" + + +def test_decrypt_json_strings_reads_values_encrypted_under_the_given_key(): + stored = encrypt_json_strings(_NESTED_PARAMS, new_encryption_key="sk-other-key") + + assert decrypt_json_strings(stored, signing_key="sk-other-key") == _NESTED_PARAMS + assert decrypt_json_strings(stored)["extra_headers"]["Authorization"] != "Bearer gateway-secret" + + +def test_encrypt_stored_json_object_rekeys_a_row_stored_as_a_json_string(): + stored_row = json.dumps(encrypt_json_strings({"api_key": "vendor-secret", "default_on": False})) + + rotated = encrypt_stored_json_object(stored_row, new_encryption_key="sk-rotated") + + assert rotated is not None + assert decrypt_if_encrypted_with(rotated["api_key"], "sk-rotated") == "vendor-secret" + assert rotated["default_on"] is False + + +@pytest.mark.parametrize("stored_row", ("not json", json.dumps(["api_key"]), 7, None)) +def test_encrypt_stored_json_object_returns_none_for_a_row_that_is_not_a_json_object(stored_row): + assert encrypt_stored_json_object(stored_row, new_encryption_key="sk-rotated") is None + + +def test_decrypt_stored_json_object_reports_a_malformed_row_without_its_contents(caplog, monkeypatch): + import logging + + from litellm._logging import verbose_proxy_logger + + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name) + + assert decrypt_stored_json_object("leaked-secret-text") == {} + assert decrypt_stored_json_object(None) == {} + + assert [record.levelno for record in caplog.records] == [logging.ERROR] + assert "leaked-secret-text" not in caplog.text diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index e379ea8212d..b336bbf38d3 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -1200,3 +1200,48 @@ def test_guardrail_from_row_keeps_a_row_without_litellm_params(guardrail_salt_ke row = _StoredGuardrailRow(guardrail_id="g-empty", guardrail_name="empty", litellm_params=None, guardrail_info=None) assert guardrail_from_row(row)["litellm_params"] is None + + +@pytest.fixture +def proxy_debug_log(caplog, monkeypatch): + import logging + + from litellm._logging import verbose_proxy_logger + + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + caplog.set_level(logging.DEBUG, logger=verbose_proxy_logger.name) + return caplog + + +def test_initialize_guardrail_keeps_the_vendor_api_key_out_of_the_logs(proxy_debug_log): + registry_module = _register_noop_initializer("quiet_init_test") + guardrail = { + "guardrail_name": "quiet-init", + "litellm_params": {"guardrail": "quiet_init_test", "mode": "pre_call", "api_key": "vendor-secret-init"}, + } + try: + assert InMemoryGuardrailHandler().initialize_guardrail(guardrail=guardrail) is not None + finally: + registry_module.guardrail_initializer_registry.pop("quiet_init_test", None) + + assert "vendor-secret-init" not in proxy_debug_log.text + + +def test_has_guardrail_params_changed_logs_only_the_changed_field_names(proxy_debug_log): + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["g-1"] = { + "guardrail_id": "g-1", + "guardrail_name": "moderation", + "litellm_params": {"guardrail": "openai_moderation", "mode": "pre_call", "api_key": "old-vendor-secret"}, + } + changed = { + "guardrail_id": "g-1", + "guardrail_name": "moderation", + "litellm_params": {"guardrail": "openai_moderation", "mode": "pre_call", "api_key": "new-vendor-secret"}, + } + + assert handler._has_guardrail_params_changed("g-1", changed) is True + + assert "api_key" in proxy_debug_log.text + assert "old-vendor-secret" not in proxy_debug_log.text + assert "new-vendor-secret" not in proxy_debug_log.text diff --git a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py index 5fa096b1523..9a5ba0e2749 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py +++ b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py @@ -595,3 +595,22 @@ async def test_migrate_covered_tables_reports_real_counts(salt_key, monkeypatch) assert by_loc["model_table"].migrated == 1 # was legacy pre, v2 post assert by_loc["model_table"].legacy == 0 # residual zero after rotation assert by_loc["model_table"].already_v2 == 1 + + +@pytest.mark.asyncio +async def test_scan_covered_tables_reports_router_settings_next_to_environment_variables(salt_key, monkeypatch): + legacy = _legacy_ct("sk-openai", monkeypatch) + config_rows = { + "environment_variables": SimpleNamespace(param_value={"OPENAI_API_KEY": legacy}), + "router_settings": SimpleNamespace(param_value={"redis_password": "plain-redis-pw", "num_retries": 2}), + } + + client = MagicMock() + _empty_covered_tables(client) + client.db.litellm_config.find_unique = AsyncMock(side_effect=lambda where: config_rows.get(where["param_name"])) + + by_loc = {r.location: r for r in await cm._scan_covered_tables(client)} + + assert by_loc["config_environment_variables"].legacy == 1 + assert by_loc["config_router_settings"].plaintext == 1 + assert by_loc["config_router_settings"].legacy == 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index f460141f581..5a08d1dd798 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -9367,6 +9367,7 @@ async def test_rotate_master_key_reencrypts_model_params_in_place( # Mock config table — no env vars mock_prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) # Mock credentials table — no credentials mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock( @@ -18537,6 +18538,7 @@ async def test_rotate_master_key_rotates_sso_identity_assertions( ) ) mock_prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock( return_value=[] ) @@ -21095,23 +21097,7 @@ class TestTeamAdminMemberKeyBudgetUpdate: assert "member_key_budgets" not in str(exc.value.detail) -@pytest.mark.asyncio -@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key") -async def test_rotate_master_key_rekeys_router_settings_and_guardrail_params(mock_rotate_mcp, monkeypatch): - import prisma - - from litellm.proxy import proxy_server - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_json_strings - from litellm.proxy.management_endpoints.key_management_endpoints import _rotate_master_key - - monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key") - monkeypatch.setattr(proxy_server, "general_settings", {}) - stored_router_settings = encrypt_json_strings({"redis_password": "redis-pw", "num_retries": 2}) - stored_guardrail_params = encrypt_json_strings( - {"guardrail": "openai_moderation", "api_key": "vendor-secret", "default_on": False} - ) - +def _rotation_prisma_client(stored_router_settings: object, stored_guardrail_params: object) -> AsyncMock: mock_prisma_client = AsyncMock() mock_prisma_client.db = MagicMock() mock_prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) @@ -21127,10 +21113,18 @@ async def test_rotate_master_key_rekeys_router_settings_and_guardrail_params(moc ) mock_prisma_client.db.litellm_guardrailstable.update = AsyncMock() mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) - mock_rotate_mcp.return_value = None + return mock_prisma_client + + +async def _rotate_to_new_master_key(mock_prisma_client: AsyncMock, monkeypatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import _rotate_master_key + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key") + monkeypatch.setattr(proxy_server, "general_settings", {}) mock_proxy_config = MagicMock() mock_proxy_config.decrypt_model_list_from_db.return_value = [] - with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config): await _rotate_master_key( prisma_client=mock_prisma_client, @@ -21141,6 +21135,26 @@ async def test_rotate_master_key_rekeys_router_settings_and_guardrail_params(moc new_master_key="sk-new-master-key", ) + +@pytest.mark.asyncio +@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key") +async def test_rotate_master_key_rekeys_router_settings_and_guardrail_params(mock_rotate_mcp, monkeypatch): + import prisma + + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_json_strings + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key") + monkeypatch.setattr(proxy_server, "general_settings", {}) + stored_router_settings = encrypt_json_strings({"redis_password": "redis-pw", "num_retries": 2}) + stored_guardrail_params = encrypt_json_strings( + {"guardrail": "openai_moderation", "api_key": "vendor-secret", "default_on": False} + ) + mock_prisma_client = _rotation_prisma_client(stored_router_settings, stored_guardrail_params) + mock_rotate_mcp.return_value = None + + await _rotate_to_new_master_key(mock_prisma_client, monkeypatch) + router_update = mock_prisma_client.db.litellm_config.update.call_args assert router_update.kwargs["where"] == {"param_name": "router_settings"} assert isinstance(router_update.kwargs["data"]["param_value"], prisma.Json) @@ -21156,3 +21170,76 @@ async def test_rotate_master_key_rekeys_router_settings_and_guardrail_params(moc assert decrypt_if_encrypted_with(rotated_guardrail_params["api_key"], "sk-new-master-key") == "vendor-secret" assert decrypt_if_encrypted_with(rotated_guardrail_params["guardrail"], "sk-new-master-key") == "openai_moderation" assert rotated_guardrail_params["default_on"] is False + + +@pytest.mark.asyncio +@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key") +async def test_rotate_master_key_rekeys_rows_stored_as_json_strings(mock_rotate_mcp, monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_json_strings + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key") + monkeypatch.setattr(proxy_server, "general_settings", {}) + stored_router_settings = json.dumps(encrypt_json_strings({"redis_password": "redis-pw"})) + stored_guardrail_params = json.dumps(encrypt_json_strings({"api_key": "vendor-secret", "default_on": False})) + mock_prisma_client = _rotation_prisma_client(stored_router_settings, stored_guardrail_params) + mock_rotate_mcp.return_value = None + + await _rotate_to_new_master_key(mock_prisma_client, monkeypatch) + + rotated_router_settings = mock_prisma_client.db.litellm_config.update.call_args.kwargs["data"]["param_value"].data + assert decrypt_if_encrypted_with(rotated_router_settings["redis_password"], "sk-new-master-key") == "redis-pw" + rotated_guardrail_params = mock_prisma_client.db.litellm_guardrailstable.update.call_args.kwargs["data"][ + "litellm_params" + ].data + assert decrypt_if_encrypted_with(rotated_guardrail_params["api_key"], "sk-new-master-key") == "vendor-secret" + assert rotated_guardrail_params["default_on"] is False + + +@pytest.mark.asyncio +@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key") +async def test_rotate_master_key_reports_rows_that_are_not_json_objects_without_their_contents( + mock_rotate_mcp, monkeypatch, caplog +): + import logging + + from litellm._logging import verbose_proxy_logger + + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name) + mock_prisma_client = _rotation_prisma_client("plaintext-router-secret", json.dumps(["vendor-secret"])) + mock_rotate_mcp.return_value = None + + await _rotate_to_new_master_key(mock_prisma_client, monkeypatch) + + skipped: Final = sorted( + record.getMessage() for record in caplog.records if record.getMessage().startswith("Master key rotation skipped") + ) + assert skipped == [ + "Master key rotation skipped config router_settings: stored value is not a JSON object", + "Master key rotation skipped guardrail guardrail-1: litellm_params is not a JSON object", + ] + assert "plaintext-router-secret" not in caplog.text + assert "vendor-secret" not in caplog.text + mock_prisma_client.db.litellm_config.update.assert_not_called() + mock_prisma_client.db.litellm_guardrailstable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_guardrail_rows_treats_only_a_missing_table_as_empty(): + from prisma.errors import TableNotFoundError + + from litellm.proxy.management_endpoints.key_management_endpoints import _guardrail_rows + + missing_table = MagicMock() + missing_table.db.litellm_guardrailstable.find_many = AsyncMock( + side_effect=TableNotFoundError( + {"user_facing_error": {"error_code": "P2021", "meta": {"table": "LiteLLM_GuardrailsTable"}}} + ) + ) + assert await _guardrail_rows(missing_table) == () + + broken_connection = MagicMock() + broken_connection.db.litellm_guardrailstable.find_many = AsyncMock(side_effect=RuntimeError("connection lost")) + with pytest.raises(RuntimeError, match="connection lost"): + await _guardrail_rows(broken_connection) diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index 306aa95a4a5..7f3474ab3b0 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -2290,3 +2290,18 @@ class TestModelRepositoryNestedEncryption: found: Final = await repo.find_by_id("gw-1") assert found is not None assert found.litellm_params == litellm_params + + @pytest.mark.asyncio + async def test_reads_use_the_same_injected_key_as_writes(self): + client: Final = MockPrismaClient() + repo: Final = ModelRepository(client, encryption_key="sk-injected-key-5678") + litellm_params: Final = {"model": "openai/gpt-5.4-mini", "extra_headers": {"Authorization": "Bearer gw"}} + + await repo.create_model(model_name="gateway-model", litellm_params=litellm_params, created_by="admin", model_id="gw-2") + + found: Final = await repo.find_by_id("gw-2") + assert found is not None + assert found.litellm_params == litellm_params + env_key_read: Final = await ModelRepository(client).find_by_id("gw-2") + assert env_key_read is not None + assert env_key_read.litellm_params["extra_headers"]["Authorization"] != "Bearer gw" From 256b64b18d5a50ddc6f4f0572407e2f9a7d3ae93 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 18:10:47 -0700 Subject: [PATCH 3/7] fix(proxy): refuse over-depth secrets on write, decode JSON-text guardrail rows in usage views, allowlist the depth-capped walkers --- .../common_utils/encrypt_decrypt_utils.py | 34 ++++++++----------- litellm/proxy/guardrails/usage_endpoints.py | 2 +- .../key_management_endpoints.py | 16 ++++++--- .../code_coverage_tests/recursive_detector.py | 2 ++ .../test_encrypt_decrypt_utils.py | 26 ++++++++++---- .../proxy/guardrails/test_usage_endpoints.py | 8 +++-- .../test_key_management_endpoints.py | 31 +++++++++++++++++ 7 files changed, 85 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 7b8a9fcb383..6310059d007 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -292,30 +292,26 @@ def _encrypted_string(value: str, new_encryption_key: str | None) -> str: def encrypt_json_strings(value: JsonValue, new_encryption_key: str | None = None, depth: int = 0) -> JsonValue: if depth > DEFAULT_MAX_RECURSE_DEPTH: - return value - match value: - case str(): - return _encrypted_string(value, new_encryption_key) - case list(): - return [encrypt_json_strings(item, new_encryption_key, depth + 1) for item in value] - case dict(): - return {key: encrypt_json_strings(item, new_encryption_key, depth + 1) for key, item in value.items()} - case _: - return value + raise ValueError(f"cannot encrypt a value nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels") + if isinstance(value, str): + return _encrypted_string(value, new_encryption_key) + if isinstance(value, list): + return [encrypt_json_strings(item, new_encryption_key, depth + 1) for item in value] + if isinstance(value, dict): + return {key: encrypt_json_strings(item, new_encryption_key, depth + 1) for key, item in value.items()} + return value def decrypt_json_strings(value: JsonValue, signing_key: str | None = None, depth: int = 0) -> JsonValue: if depth > DEFAULT_MAX_RECURSE_DEPTH: return value - match value: - case str(): - return _plaintext_of(value, signing_key) - case list(): - return [decrypt_json_strings(item, signing_key, depth + 1) for item in value] - case dict(): - return {key: decrypt_json_strings(item, signing_key, depth + 1) for key, item in value.items()} - case _: - return value + if isinstance(value, str): + return _plaintext_of(value, signing_key) + if isinstance(value, list): + return [decrypt_json_strings(item, signing_key, depth + 1) for item in value] + if isinstance(value, dict): + return {key: decrypt_json_strings(item, signing_key, depth + 1) for key, item in value.items()} + return value def parse_stored_json_object(value: object) -> dict[str, JsonValue] | None: diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 21cfe382fad..397a405a775 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -432,7 +432,7 @@ def _to_dict(value: object) -> dict[str, Any]: def _guardrail_litellm_params(g: "_DbOrConfigGuardrail") -> Mapping[str, object]: value: Final = _get_guardrail_field(g, "litellm_params") - if isinstance(value, dict): + if isinstance(value, (dict, str)): return decrypt_stored_json_object(value) return _to_dict(value) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4132ddb708f..25cfcd2efbf 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5243,11 +5243,13 @@ async def _rekey_config_section( "Master key rotation skipped config %s: stored value is not a JSON object", param_name ) return + try: + rekeyed: Final = encrypt_config_section(param_name, section, new_encryption_key=new_master_key) + except ValueError as too_deep: + verbose_proxy_logger.error("Master key rotation skipped config %s: %s", param_name, too_deep) + return await _config_table(prisma_client).update( - where={"param_name": param_name}, - data={ - "param_value": prisma.Json(encrypt_config_section(param_name, section, new_encryption_key=new_master_key)) - }, + where={"param_name": param_name}, data={"param_value": prisma.Json(rekeyed)} ) @@ -5256,7 +5258,11 @@ async def _rekey_guardrail_row( ) -> None: import prisma - rekeyed: Final = encrypt_stored_json_object(guardrail_row.litellm_params, new_encryption_key=new_master_key) + try: + rekeyed: Final = encrypt_stored_json_object(guardrail_row.litellm_params, new_encryption_key=new_master_key) + except ValueError as too_deep: + verbose_proxy_logger.error("Master key rotation skipped guardrail %s: %s", guardrail_row.guardrail_id, too_deep) + return if rekeyed is None: verbose_proxy_logger.error( "Master key rotation skipped guardrail %s: litellm_params is not a JSON object", guardrail_row.guardrail_id diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 659dc438f2d..b5105d78cd1 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -73,6 +73,8 @@ IGNORE_FUNCTIONS = [ "_replace_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). "_sort_processed_sets", # bounded by the nesting depth of the log-record extra it walks (a finite JSON tree, no cycles possible). "scrub_json_strings", # max depth set (MAX_SCRUB_DEPTH); fails closed by returning "[Filtered]" for anything nested past the cap. + "encrypt_json_strings", # max depth set. + "decrypt_json_strings", # max depth set. ] diff --git a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py index 6856ff46a44..92f6c462876 100644 --- a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -308,18 +308,28 @@ def test_encrypt_json_strings_without_any_key_stores_the_value_as_sent(monkeypat assert decrypt_json_strings(_NESTED_PARAMS) == _NESTED_PARAMS -def test_encrypt_json_strings_stops_descending_past_the_recursion_cap(): - leaf = {"secret": "deep-secret"} - value = leaf +def test_encrypt_json_strings_refuses_a_value_nested_past_the_recursion_cap(): + value = {"secret": "deep-secret"} for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1): value = {"child": value} + with pytest.raises(ValueError, match=f"nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels"): + encrypt_json_strings(value) + assert decrypt_json_strings(value) == value + + +def test_encrypt_json_strings_encrypts_a_value_nested_at_the_recursion_cap(): + value = "deep-secret" + for _ in range(DEFAULT_MAX_RECURSE_DEPTH): + value = {"child": value} + stored = encrypt_json_strings(value) deepest = stored - for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1): + for _ in range(DEFAULT_MAX_RECURSE_DEPTH): deepest = deepest["child"] - assert deepest == leaf + assert deepest != "deep-secret" + assert decrypt_json_strings(stored) == value def test_json_value_coerces_non_json_containers(): @@ -339,7 +349,11 @@ def test_decrypt_stored_json_object_reads_json_text_and_parsed_rows_alike(): def test_encrypt_config_section_encrypts_router_settings_string_leaves_only(): - router_settings = {"redis_password": "redis-pw", "num_retries": 2, "fallbacks": [{"gpt-5.5-mini": ["gpt-5.4-mini"]}]} + router_settings = { + "redis_password": "redis-pw", + "num_retries": 2, + "fallbacks": [{"gpt-5.5-mini": ["gpt-5.4-mini"]}], + } stored = encrypt_config_section("router_settings", router_settings) diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index b5588dd6e5e..b40c510a3f2 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -185,7 +185,8 @@ async def test_detail_db_row_still_resolves(): @pytest.mark.asyncio -async def test_detail_encrypted_db_row_resolves_its_provider(monkeypatch): +@pytest.mark.parametrize("stored_as_json_text", [False, True]) +async def test_detail_encrypted_db_row_resolves_its_provider(monkeypatch, stored_as_json_text: bool): import json from litellm.proxy import proxy_server @@ -194,8 +195,9 @@ async def test_detail_encrypted_db_row_resolves_its_provider(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "sk-usage-salt-1234") monkeypatch.setattr(proxy_server, "general_settings", {}) row = _db_row(guardrail_id="db-encrypted", provider="aim") - row.litellm_params = json.loads(encrypted_guardrail_litellm_params({"guardrail": "aim", "mode": "pre_call"})) - assert row.litellm_params["guardrail"] != "aim" + stored = encrypted_guardrail_litellm_params({"guardrail": "aim", "mode": "pre_call"}) + assert json.loads(stored)["guardrail"] != "aim" + row.litellm_params = stored if stored_as_json_text else json.loads(stored) prisma = _prisma(find_unique=row) handler = _config_handler() p1, p2 = _patches(prisma, handler) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 5a08d1dd798..b7c21e02ad4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -21225,6 +21225,37 @@ async def test_rotate_master_key_reports_rows_that_are_not_json_objects_without_ mock_prisma_client.db.litellm_guardrailstable.update.assert_not_called() +@pytest.mark.asyncio +@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key") +async def test_rotate_master_key_skips_rows_nested_past_the_recursion_cap(mock_rotate_mcp, monkeypatch, caplog): + import logging + + from litellm._logging import verbose_proxy_logger + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name) + too_deep: dict[str, object] = {"secret": "vendor-secret"} + for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1): + too_deep = {"child": too_deep} + mock_prisma_client = _rotation_prisma_client(too_deep, too_deep) + mock_rotate_mcp.return_value = None + + await _rotate_to_new_master_key(mock_prisma_client, monkeypatch) + + skipped: Final = sorted( + record.getMessage() for record in caplog.records if record.getMessage().startswith("Master key rotation skipped") + ) + reason: Final = f"cannot encrypt a value nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels" + assert skipped == [ + f"Master key rotation skipped config router_settings: {reason}", + f"Master key rotation skipped guardrail guardrail-1: {reason}", + ] + assert "vendor-secret" not in caplog.text + mock_prisma_client.db.litellm_config.update.assert_not_called() + mock_prisma_client.db.litellm_guardrailstable.update.assert_not_called() + + @pytest.mark.asyncio async def test_guardrail_rows_treats_only_a_missing_table_as_empty(): from prisma.errors import TableNotFoundError From 758f162e47cd50eb89cb2fe94d0bd0088217d119 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 19:50:36 -0700 Subject: [PATCH 4/7] fix(proxy): share the model params encryption rule with the tag writer and decrypt stored session tags before the admin check --- .../model_management_endpoints.py | 50 ++++------- .../tag_management_endpoints.py | 30 +++---- litellm/repositories/model_repository.py | 33 ++++++- .../test_model_management_endpoints.py | 26 ++++++ .../test_tag_management_endpoints.py | 87 ++++++++----------- tests/unit/repositories/test_repositories.py | 24 +++++ 6 files changed, 143 insertions(+), 107 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index e9d85060599..0e0163e8d2e 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -22,7 +22,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator import litellm from litellm._logging import verbose_proxy_logger @@ -63,8 +63,8 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( publish_config_change, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_json_strings, decrypt_value_helper, - encrypt_json_strings, json_value, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -98,7 +98,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.credentials_repository import CredentialsRepository -from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.model_repository import ModelRepository, encrypt_model_litellm_params from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ModelTableRepository from litellm.repositories.team_repository import TeamRepository @@ -437,33 +437,6 @@ def _effective_complexity_router_config( } -def _encrypted_complexity_router_config(config: object, new_encryption_key: str | None = None) -> JsonValue: - value: Final = json_value(config) - if not isinstance(value, dict): - return value - jev: Final = value.get("jev_classifier_config") - if not isinstance(jev, dict): - return value - api_key: Final = jev.get("api_key") - if not isinstance(api_key, str): - return value - encrypted_api_key: Final = encrypt_json_strings(api_key, new_encryption_key=new_encryption_key) - return {**value, "jev_classifier_config": {**jev, "api_key": encrypted_api_key}} - - -def _encrypted_litellm_params( - litellm_params: Mapping[str, object], new_encryption_key: str | None = None -) -> dict[str, JsonValue]: - return { - key: ( - _encrypted_complexity_router_config(value, new_encryption_key) - if key == "complexity_router_config" - else encrypt_json_strings(json_value(value), new_encryption_key=new_encryption_key) - ) - for key, value in litellm_params.items() - } - - def _effective_model( incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None ) -> str | None: @@ -1023,7 +996,7 @@ def update_db_model( ) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() } - merged_litellm_params.update(_encrypted_litellm_params({**merged_litellm_params, **patched_params})) + merged_litellm_params.update(encrypt_model_litellm_params({**merged_litellm_params, **patched_params})) # update model info if updated_patch.model_info: @@ -1527,7 +1500,7 @@ async def _add_model_to_db( # encrypt litellm params # _litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) _original_litellm_model_name: Final = model_params.litellm_params.model - for k, v in _encrypted_litellm_params(_litellm_params_dict, new_encryption_key).items(): + for k, v in encrypt_model_litellm_params(_litellm_params_dict, new_encryption_key).items(): model_params.litellm_params[k] = v _data: Final[dict] = { "model_id": model_params.model_info.id, @@ -2009,6 +1982,15 @@ def _canonical_session_tags(tags: Sequence[AwsSessionTag]) -> tuple[tuple[str, s return tuple(sorted((tag["Key"], tag["Value"]) for tag in tags)) +_SESSION_TAGS: Final = TypeAdapter(list[AwsSessionTag]) + + +def _stored_session_tags(existing_litellm_params: GenericLiteLLMParams | None) -> Sequence[AwsSessionTag] | None: + if existing_litellm_params is None or existing_litellm_params.aws_session_tags is None: + return None + return _SESSION_TAGS.validate_python(decrypt_json_strings(json_value(existing_litellm_params.aws_session_tags))) + + class ModelManagementAuthChecks: """ Common auth checks for model management endpoints @@ -2074,7 +2056,7 @@ class ModelManagementAuthChecks: return True if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return True - existing_tags: Final = existing_litellm_params.aws_session_tags if existing_litellm_params is not None else None + existing_tags: Final = _stored_session_tags(existing_litellm_params) if existing_tags is not None and _canonical_session_tags(existing_tags) == _canonical_session_tags( litellm_params.aws_session_tags ): @@ -2695,7 +2677,7 @@ async def update_model( ) _mp: Final[dict[str, object]] = model_params.litellm_params.dict() - merged_dictionary: Final = _encrypted_litellm_params( + merged_dictionary: Final = encrypt_model_litellm_params( { key: _existing_litellm_params_dict[key] if value is None else requested_params[key] for key, value in _mp.items() diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index ab33d4bd766..40f9273341e 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -21,6 +21,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, parse_stored_json_object from litellm.proxy.common_utils.user_api_key_cache import ( tag_cache_key, tag_registry_cache_key, @@ -30,7 +31,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity, ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity -from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.model_repository import ModelRepository, encrypt_model_litellm_params from litellm.repositories.table_repositories import ( DailyTagSpendRepository, TagRepository, @@ -358,7 +359,6 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): raise HTTPException(status_code=500, detail="Database not connected") try: - # Get current model from database to preserve encrypted fields db_model = await ModelRepository(prisma_client).table.find_unique(where={"model_id": deployment.model_info.id}) if db_model is None: @@ -367,24 +367,20 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): detail=f"Model {deployment.model_info.id} not found in database", ) - # Prisma returns litellm_params as dict (already parsed from JSON) - existing_params = db_model.litellm_params - if isinstance(existing_params, str): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub is str - # If it's a string, parse it - existing_params = json.loads(existing_params) - elif not isinstance(existing_params, dict): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub - raise Exception(f"Unexpected litellm_params type: {type(existing_params)}") + stored_params: Final = parse_stored_json_object(db_model.litellm_params) + if stored_params is None: + raise Exception(f"Unexpected litellm_params type: {type(db_model.litellm_params)}") + existing_params: Final = {key: decrypt_json_strings(value) for key, value in stored_params.items()} + stored_tags: Final = existing_params.get("tags") + tags: Final = stored_tags if isinstance(stored_tags, list) else [] + if tag in tags: + return - # Add tag to tags array (preserve encryption of other fields) - if "tags" not in existing_params: - existing_params["tags"] = [] - if tag not in existing_params["tags"]: - existing_params["tags"].append(tag) - - # Update database with modified params (keeps encrypted fields encrypted) await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_info.id}, - data={"litellm_params": json.dumps(existing_params)}, + data={ + "litellm_params": json.dumps(encrypt_model_litellm_params({**existing_params, "tags": [*tags, tag]})) + }, ) except Exception as e: verbose_proxy_logger.exception("Error adding tag to deployment: %s", e) diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 1371a2b1a25..790f83fc30a 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -6,6 +6,8 @@ import json from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final +from pydantic import JsonValue + from litellm.models.model import LiteLLM_ProxyModelTable from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_json_strings, @@ -24,6 +26,32 @@ class _ProxyModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_Pr table_name = "litellm_proxymodeltable" +def _complexity_router_config_with_encrypted_api_key(config: JsonValue, new_encryption_key: str | None) -> JsonValue: + if not isinstance(config, dict): + return config + jev: Final = config.get("jev_classifier_config") + if not isinstance(jev, dict): + return config + api_key: Final = jev.get("api_key") + if not isinstance(api_key, str): + return config + encrypted_api_key: Final = encrypt_json_strings(api_key, new_encryption_key=new_encryption_key) + return {**config, "jev_classifier_config": {**jev, "api_key": encrypted_api_key}} + + +def encrypt_model_litellm_params( + litellm_params: Mapping[str, object], new_encryption_key: str | None = None +) -> dict[str, JsonValue]: + return { + key: ( + _complexity_router_config_with_encrypted_api_key(json_value(value), new_encryption_key) + if key == "complexity_router_config" + else encrypt_json_strings(json_value(value), new_encryption_key=new_encryption_key) + ) + for key, value in litellm_params.items() + } + + class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): """Repository for proxy model database operations with encryption support.""" @@ -40,10 +68,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): return LiteLLM_ProxyModelTable def _encrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]: - return { - key: encrypt_json_strings(json_value(value), new_encryption_key=self._encryption_key) - for key, value in litellm_params.items() - } + return encrypt_model_litellm_params(litellm_params, new_encryption_key=self._encryption_key) def _decrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]: return { diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 289d1438467..7303c2eda00 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -507,6 +507,32 @@ class TestModelManagementAuthChecks: ) assert exc_info.value.code == "403" + def test_can_user_set_aws_session_tags_compares_against_the_decrypted_stored_tags(self, monkeypatch): + from litellm.repositories.model_repository import encrypt_model_litellm_params + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-session-tags-salt-8627") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + stored: Final = encrypt_model_litellm_params({"aws_session_tags": [{"Key": "team", "Value": "genai"}]}) + assert stored["aws_session_tags"] != [{"Key": "team", "Value": "genai"}] + existing: Final = LiteLLM_Params(model="bedrock/test_model", aws_session_tags=stored["aws_session_tags"]) + + result = ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=LiteLLM_Params( + model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "genai"}] + ), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=existing, + ) + assert result is True + with pytest.raises(Exception, match="Only a proxy admin can set aws_session_tags"): + ModelManagementAuthChecks.can_user_set_aws_session_tags( + litellm_params=LiteLLM_Params( + model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "platform"}] + ), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=existing, + ) + @pytest.mark.asyncio async def test_add_new_model_rejects_aws_session_tags_for_non_admin(self): from litellm.proxy._types import ProxyException diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index 3cfdd345a45..2bc57b1b58f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -1428,58 +1428,6 @@ async def test_add_tag_to_deployment_with_string_params(): assert updated_params["api_key"] == "encrypted_claude_key" -@pytest.mark.asyncio -async def test_add_tag_to_deployment_no_duplicate_tags(): - """ - Test that _add_tag_to_deployment doesn't add duplicate tags - """ - from unittest.mock import AsyncMock, Mock - - from litellm.proxy.management_endpoints.tag_management_endpoints import ( - _add_tag_to_deployment, - ) - from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo - - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - # Setup prisma mocks - mock_db = Mock() - mock_prisma.db = mock_db - - # Mock the database model with existing tags - db_model = Mock() - db_model.model_id = "model-789" - db_model.litellm_params = { - "model": "gpt-4", - "api_key": "encrypted_key", - "tags": ["existing-tag", "another-tag"], - } - - # Mock find_unique to return the db model - mock_db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_model) - - # Mock update - mock_db.litellm_proxymodeltable.update = AsyncMock(return_value=db_model) - - # Create deployment - deployment = Deployment( - model_name="gpt-4", - litellm_params=LiteLLM_Params(model="gpt-4"), - model_info=ModelInfo(id="model-789"), - ) - - # Try to add an existing tag - await _add_tag_to_deployment(deployment, "existing-tag") - - # Verify update was called - update_call = mock_db.litellm_proxymodeltable.update.call_args - updated_params = json.loads(update_call[1]["data"]["litellm_params"]) - - # Verify no duplicate tags - assert updated_params["tags"].count("existing-tag") == 1 - assert len(updated_params["tags"]) == 2 - assert "another-tag" in updated_params["tags"] - - @pytest.mark.asyncio async def test_add_tag_to_deployment_model_not_found(): """ @@ -1513,3 +1461,38 @@ async def test_add_tag_to_deployment_model_not_found(): assert exc_info.value.status_code == 500 assert "not found in database" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_add_tag_to_deployment_merges_into_the_decrypted_tags_and_writes_ciphertext(monkeypatch): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, encrypt_json_strings + from litellm.proxy.management_endpoints.tag_management_endpoints import _add_tag_to_deployment + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-tag-salt-8627") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + plaintext_params = {"model": "openai/gpt-5.4-mini", "api_key": "sk-placeholder", "tags": ["existing-tag"]} + stored_params = encrypt_json_strings(plaintext_params) + assert stored_params != plaintext_params + deployment = Deployment( + model_name="gpt-5.4-mini", + litellm_params=LiteLLM_Params(model="openai/gpt-5.4-mini"), + model_info=ModelInfo(id="model-8627"), + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + mock_db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=SimpleNamespace(model_id="model-8627", litellm_params=stored_params) + ) + mock_db.litellm_proxymodeltable.update = AsyncMock() + + await _add_tag_to_deployment(deployment, "existing-tag") + mock_db.litellm_proxymodeltable.update.assert_not_awaited() + + await _add_tag_to_deployment(deployment, "new-tag") + written = json.loads(mock_db.litellm_proxymodeltable.update.call_args.kwargs["data"]["litellm_params"]) + assert "new-tag" not in written["tags"] + assert written["api_key"] != "sk-placeholder" + assert decrypt_json_strings(written) == {**plaintext_params, "tags": ["existing-tag", "new-tag"]} diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index 7f3474ab3b0..b8fa6c18058 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -2291,6 +2291,30 @@ class TestModelRepositoryNestedEncryption: assert found is not None assert found.litellm_params == litellm_params + @pytest.mark.asyncio + async def test_create_model_keeps_complexity_router_classifier_fields_readable_in_sql(self): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with + + client: Final = MockPrismaClient() + repo: Final = ModelRepository(client, encryption_key=self.SALT_KEY) + complexity_router_config: Final = { + "classifier_type": "jev", + "jev_classifier_config": {"api_base": "https://classifier.example", "api_key": "sk-classifier-secret"}, + } + litellm_params: Final = {"model": "auto_router/tuned", "complexity_router_config": complexity_router_config} + + await repo.create_model(model_name="auto-router", litellm_params=litellm_params, created_by="admin", model_id="ar-1") + + stored_text: Final = client.db.litellm_proxymodeltable._records["ar-1"]["litellm_params"] + assert "sk-classifier-secret" not in stored_text + stored_config: Final = json.loads(stored_text)["complexity_router_config"] + assert stored_config["classifier_type"] == "jev" + assert stored_config["jev_classifier_config"]["api_base"] == "https://classifier.example" + assert decrypt_if_encrypted_with(stored_config["jev_classifier_config"]["api_key"], self.SALT_KEY) == "sk-classifier-secret" + found: Final = await repo.find_by_id("ar-1") + assert found is not None + assert found.litellm_params == litellm_params + @pytest.mark.asyncio async def test_reads_use_the_same_injected_key_as_writes(self): client: Final = MockPrismaClient() From bdc04a506059c13adfc11ed5e631bcdf4141a01c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 20:33:59 -0700 Subject: [PATCH 5/7] fix(proxy): encrypt the jev classifier api_base and decrypt stored classifier transport before the patch merge compares it --- .../model_management_endpoints.py | 2 +- litellm/repositories/model_repository.py | 23 +++++--- .../test_encrypt_decrypt_utils.py | 32 ++++++++--- .../test_model_management_endpoints.py | 57 ++++++++++--------- tests/unit/repositories/test_repositories.py | 15 +++-- 5 files changed, 83 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0e0163e8d2e..468153aed49 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -419,7 +419,7 @@ def _effective_complexity_router_config( if not isinstance(incoming_jev, Mapping) or not isinstance(existing_jev, Mapping): return incoming supplied: Final = TypeAdapter(dict[str, object]).validate_python(incoming_jev) - stored: Final = TypeAdapter(dict[str, object]).validate_python(existing_jev) + stored: Final = TypeAdapter(dict[str, object]).validate_python(decrypt_json_strings(json_value(existing_jev))) same_base: Final = "api_base" not in supplied or supplied["api_base"] == stored.get("api_base") transport: Final = MappingProxyType( { diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 790f83fc30a..13aea6c18ea 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -26,17 +26,26 @@ class _ProxyModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_Pr table_name = "litellm_proxymodeltable" -def _complexity_router_config_with_encrypted_api_key(config: JsonValue, new_encryption_key: str | None) -> JsonValue: +JEV_CLASSIFIER_CREDENTIAL_FIELDS: Final = frozenset({"api_key", "api_base"}) + + +def _complexity_router_config_with_encrypted_classifier_credentials( + config: JsonValue, new_encryption_key: str | None +) -> JsonValue: if not isinstance(config, dict): return config jev: Final = config.get("jev_classifier_config") if not isinstance(jev, dict): return config - api_key: Final = jev.get("api_key") - if not isinstance(api_key, str): - return config - encrypted_api_key: Final = encrypt_json_strings(api_key, new_encryption_key=new_encryption_key) - return {**config, "jev_classifier_config": {**jev, "api_key": encrypted_api_key}} + encrypted_jev: Final = { + key: ( + encrypt_json_strings(value, new_encryption_key=new_encryption_key) + if key in JEV_CLASSIFIER_CREDENTIAL_FIELDS and isinstance(value, str) + else value + ) + for key, value in jev.items() + } + return {**config, "jev_classifier_config": encrypted_jev} def encrypt_model_litellm_params( @@ -44,7 +53,7 @@ def encrypt_model_litellm_params( ) -> dict[str, JsonValue]: return { key: ( - _complexity_router_config_with_encrypted_api_key(json_value(value), new_encryption_key) + _complexity_router_config_with_encrypted_classifier_credentials(json_value(value), new_encryption_key) if key == "complexity_router_config" else encrypt_json_strings(json_value(value), new_encryption_key=new_encryption_key) ) diff --git a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py index 92f6c462876..4f80f244b55 100644 --- a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -9,8 +9,11 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130 import base64 import json from types import MappingProxyType +from typing import Final import pytest +from pydantic import JsonValue, TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.proxy import proxy_server @@ -260,13 +263,26 @@ _NESTED_PARAMS = { } -def _nested_string_leaves(stored: dict) -> list: - return [ - stored["model"], - stored["extra_headers"]["Authorization"], - stored["extra_headers"]["X-Trace"], - stored["fallbacks"][0]["gpt-5.5-mini"][0], - ] +class _StoredNestedParams(TypedDict): + model: ReadOnly[str] + extra_headers: ReadOnly[dict[str, str]] + fallbacks: ReadOnly[list[dict[str, list[str]]]] + rpm: ReadOnly[int] + enabled: ReadOnly[bool] + api_base: ReadOnly[None] + + +_STORED_NESTED_PARAMS: Final = TypeAdapter(_StoredNestedParams) + + +def _nested_string_leaves(stored: JsonValue) -> tuple[str, str, str, str]: + params: Final = _STORED_NESTED_PARAMS.validate_python(stored) + return ( + params["model"], + params["extra_headers"]["Authorization"], + params["extra_headers"]["X-Trace"], + params["fallbacks"][0]["gpt-5.5-mini"][0], + ) @pytest.mark.parametrize("use_aes", [False, True]) @@ -279,7 +295,7 @@ def test_encrypt_json_strings_encrypts_every_nested_string_leaf(monkeypatch, use assert (stored["rpm"], stored["enabled"], stored["api_base"]) == (10, True, None) plaintexts = _nested_string_leaves(_NESTED_PARAMS) assert all(leaf != plain for leaf, plain in zip(_nested_string_leaves(stored), plaintexts)) - assert [decrypt_if_encrypted_with(leaf, SALT_KEY) for leaf in _nested_string_leaves(stored)] == plaintexts + assert tuple(decrypt_if_encrypted_with(leaf, SALT_KEY) for leaf in _nested_string_leaves(stored)) == plaintexts assert decrypt_json_strings(stored) == _NESTED_PARAMS diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 7303c2eda00..90b0420ab8e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -23,7 +23,11 @@ from litellm.proxy._types import ( ReconcileOutcome, UserAPIKeyAuth, ) -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_json_strings, + encrypt_value_helper, + json_value, +) from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, _get_team_deployments, @@ -7570,39 +7574,38 @@ class TestTeamMemberAutoRouterWrites: assert saved_info["team_id"] == "member-team" assert saved_info["access_groups"] == ["retained-admin-group"] - @staticmethod - def _with_decrypted_jev_key(config: Mapping[str, object]) -> dict[str, object]: - jev: Final = config.get("jev_classifier_config") - if not isinstance(jev, dict) or not isinstance(jev.get("api_key"), str): - return dict(config) - decrypted_key: Final = decrypt_value_helper(jev["api_key"], key="api_key") - return {**config, "jev_classifier_config": {**jev, "api_key": decrypted_key}} - @pytest.mark.asyncio @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) - @pytest.mark.parametrize("change", ["save", "rotate", "move", "move-without-key", "reset", "heuristic"]) - async def test_jev_dashboard_save_preserves_server_transport(self, endpoint: str, change: str) -> None: + @pytest.mark.parametrize("stored_as", ["plaintext", "ciphertext"]) + @pytest.mark.parametrize( + "change", ["save", "rotate", "move", "move-without-key", "same-base", "reset", "heuristic"] + ) + async def test_jev_dashboard_save_preserves_server_transport( + self, endpoint: str, stored_as: str, change: str + ) -> None: + from litellm.repositories.model_repository import encrypt_model_litellm_params + original: Final = self._row() transport: Final = {"api_key": "synthetic-original-jev-key", "api_base": "https://jev.example.com"} - stored_config: Final = { - "classifier_type": "jev", - "tiers": {"SIMPLE": "allowed"}, - "jev_classifier_config": {**transport, "instructions": "Old instructions", "timeout_ms": 6100}, + plaintext_params: Final = { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "jev", + "tiers": {"SIMPLE": "allowed"}, + "jev_classifier_config": {**transport, "instructions": "Old instructions", "timeout_ms": 6100}, + }, } - row: Final = original.model_copy( - update={ - "litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": stored_config, - }, - } + stored_params: Final = ( + plaintext_params if stored_as == "plaintext" else encrypt_model_litellm_params(plaintext_params) ) + row: Final = original.model_copy(update={"litellm_params": stored_params}) database: Final = self._database(self._team(), row) overrides: Final = { "save": {}, "rotate": {"api_key": "synthetic-replacement-jev-key"}, "move": {"api_base": "https://new-jev.example.com", "api_key": "synthetic-replacement-jev-key"}, "move-without-key": {"api_base": "https://new-jev.example.com"}, + "same-base": {"api_base": "https://jev.example.com"}, "reset": {"api_key": None, "api_base": None}, "heuristic": {}, }[change] @@ -7634,10 +7637,12 @@ class TestTeamMemberAutoRouterWrites: else {**config, "jev_classifier_config": {**transport, "timeout_ms": 8100, **overrides}} ) stored_jev: Final = saved.get("jev_classifier_config") - if isinstance(stored_jev, dict) and isinstance(stored_jev.get("api_key"), str): - assert stored_jev["api_key"] != expected["jev_classifier_config"]["api_key"] - assert self._with_decrypted_jev_key(saved) == expected - assert row.litellm_params["complexity_router_config"] == stored_config + if isinstance(stored_jev, dict): + for field in ("api_key", "api_base"): + if isinstance(stored_jev.get(field), str): + assert stored_jev[field] != expected["jev_classifier_config"][field] + assert decrypt_json_strings(json_value(saved)) == expected + assert row.litellm_params == stored_params assert request.litellm_params.complexity_router_config == config @pytest.mark.asyncio diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index b8fa6c18058..fd53072582a 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -2292,14 +2292,18 @@ class TestModelRepositoryNestedEncryption: assert found.litellm_params == litellm_params @pytest.mark.asyncio - async def test_create_model_keeps_complexity_router_classifier_fields_readable_in_sql(self): + async def test_create_model_encrypts_classifier_credentials_and_keeps_classifier_fields_readable_in_sql(self): from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with client: Final = MockPrismaClient() repo: Final = ModelRepository(client, encryption_key=self.SALT_KEY) complexity_router_config: Final = { "classifier_type": "jev", - "jev_classifier_config": {"api_base": "https://classifier.example", "api_key": "sk-classifier-secret"}, + "jev_classifier_config": { + "api_base": "https://classifier.example", + "api_key": "sk-classifier-secret", + "instructions": "Rate the question", + }, } litellm_params: Final = {"model": "auto_router/tuned", "complexity_router_config": complexity_router_config} @@ -2307,10 +2311,13 @@ class TestModelRepositoryNestedEncryption: stored_text: Final = client.db.litellm_proxymodeltable._records["ar-1"]["litellm_params"] assert "sk-classifier-secret" not in stored_text + assert "classifier.example" not in stored_text stored_config: Final = json.loads(stored_text)["complexity_router_config"] assert stored_config["classifier_type"] == "jev" - assert stored_config["jev_classifier_config"]["api_base"] == "https://classifier.example" - assert decrypt_if_encrypted_with(stored_config["jev_classifier_config"]["api_key"], self.SALT_KEY) == "sk-classifier-secret" + assert stored_config["jev_classifier_config"]["instructions"] == "Rate the question" + stored_jev: Final = stored_config["jev_classifier_config"] + assert decrypt_if_encrypted_with(stored_jev["api_base"], self.SALT_KEY) == "https://classifier.example" + assert decrypt_if_encrypted_with(stored_jev["api_key"], self.SALT_KEY) == "sk-classifier-secret" found: Final = await repo.find_by_id("ar-1") assert found is not None assert found.litellm_params == litellm_params From a1089f1f70a0ec9adb50950974510873c9908bdb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:31:28 -0700 Subject: [PATCH 6/7] test: decrypt the stored classifier config and drop the known master key literal after the main merge --- .../management_endpoints/test_key_management_endpoints.py | 2 +- .../proxy/management_helpers/test_auto_router_permissions.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 3fe92d723e7..2027c2daf1b 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -22370,7 +22370,7 @@ async def _rotate_to_new_master_key(mock_prisma_client: AsyncMock, monkeypatch) await rotate_master_key( prisma_client=mock_prisma_client, user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="test-user" + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="test-user" ), current_master_key="sk-old-master-key", new_master_key="sk-new-master-key", diff --git a/tests/unit/proxy/management_helpers/test_auto_router_permissions.py b/tests/unit/proxy/management_helpers/test_auto_router_permissions.py index e161ea3ecaf..35fddd25146 100644 --- a/tests/unit/proxy/management_helpers/test_auto_router_permissions.py +++ b/tests/unit/proxy/management_helpers/test_auto_router_permissions.py @@ -17,6 +17,7 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings from litellm.proxy.management_helpers.auto_router_permissions import ( MemberAutoRouterDependencyObjects, authorize_member_auto_router_dependencies, @@ -316,7 +317,7 @@ async def test_member_authorization_and_persistence_resolve_the_same_classifier( grant: Final = await operation persisted: Final = update_db_model(existing, patch) saved: Final = RequestComplexityRouterConfig.model_validate( - json.loads(persisted["litellm_params"])["complexity_router_config"] + decrypt_json_strings(json.loads(persisted["litellm_params"])["complexity_router_config"]) ) assert grant.config == saved assert saved.jev_classifier_config is not None From 18a28f745782fb8a3987b451a02e84e1e20356a8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:56:55 -0800 Subject: [PATCH 7/7] fix(auto-router): decrypt the stored classifier connection before checking member edits --- .../auto_router_permissions.py | 16 ++++++---- .../test_model_management_endpoints.py | 31 +++++++++++++++++++ 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py index e6a79229ba8..b7ff534f468 100644 --- a/litellm/proxy/management_helpers/auto_router_permissions.py +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -28,7 +28,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo check_team_member_model_access, # pyright: ignore[reportPrivateUsage] # shared membership authorization owner ) from litellm.proxy.auth.team_grants import team_model_aliases -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, decrypt_value_helper, json_value from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import DatabaseClient from litellm.repositories.project_repository import ProjectRepository @@ -130,6 +130,14 @@ class _RouterConfigSource(LiteLLMBaseModel): complexity_router_config: Mapping[str, object] | None = None +_ROUTER_CONFIG: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + + +def _decrypted_router_config(existing: Deployment) -> Mapping[str, object] | None: + stored: Final = _RouterConfigSource.model_validate(existing.litellm_params.model_dump()).complexity_router_config + return _ROUTER_CONFIG.validate_python(decrypt_json_strings(json_value(stored))) + + class _MembershipKey(TypedDict): user_id: ReadOnly[str] team_id: ReadOnly[str] @@ -380,11 +388,7 @@ async def authorize_member_auto_router_write( if existing is not None and incoming.model_name not in (None, public_name, existing.model_name): raise HTTPException(status_code=403, detail="Team members cannot rename an auto router.") supplied_config: Final = _RouterConfigSource.model_validate(params.model_dump()).complexity_router_config - stored_config: Final = ( - _RouterConfigSource.model_validate(existing.litellm_params.model_dump()).complexity_router_config - if existing is not None - else None - ) + stored_config: Final = _decrypted_router_config(existing) if existing is not None else None resolved_config: Final = resolve_complexity_router_config_write(supplied_config, stored_config) if resolved_config.supplied_connection_fields: raise HTTPException(status_code=403, detail="Team members cannot change classifier connections.") diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 7fb88dfa1ac..487ce5d7427 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -9418,6 +9418,37 @@ class TestTeamMemberAutoRouterWrites: } assert "retained-laya-secret" in row.model_dump_json() + @pytest.mark.asyncio + @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) + async def test_member_save_over_an_encrypted_stored_connection_keeps_it(self, endpoint: str) -> None: + from litellm.repositories.model_repository import encrypt_model_litellm_params + + config: Final = { + "classifier_type": "jev", "tiers": {"SIMPLE": "allowed"}, + "jev_classifier_config": {"provider": "laya", "model": "english"}, + } + connection: Final = {"api_key": "retained-laya-secret", "api_base": "https://laya.test"} + stored_params: Final = encrypt_model_litellm_params({ + "model": "auto_router/complexity_router", + "complexity_router_config": {**config, "jev_classifier_config": {**config["jev_classifier_config"], **connection}}, + }) + row: Final = self._row().model_copy(update={"litellm_params": dict(stored_params)}) + team: Final = self._team().model_copy(update={"models": ["allowed", "laya/english"]}) + database: Final = self._database(team, row) + database.transaction.litellm_proxymodeltable.update.return_value = row + request: Final = updateDeployment( + litellm_params=updateLiteLLMParams.model_validate({"complexity_router_config": config}), + model_info=ModelInfo(id=row.model_id, team_id="member-team"), + ) + actor: Final = UserAPIKeyAuth( + user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER, models=["allowed", "laya/english"], config={"timeout": 60}, + ) + with self._environment(database, row): + await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor)) + written: Final = database.transaction.litellm_proxymodeltable.update.await_args.kwargs["data"] + saved: Final = decrypt_json_strings(json_value(json.loads(written["litellm_params"])["complexity_router_config"])) + assert saved["opensource_classifier_config"] == {**config["jev_classifier_config"], **connection} + @pytest.mark.asyncio @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) @pytest.mark.parametrize("access", ["owner", "peer", "limited-key"])