diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index e5f6bc5dda8..f9f47c0562d 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -3,9 +3,12 @@ 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 typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 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)). @@ -296,3 +299,99 @@ 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]] = TypeAdapter(JsonValue) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _is_mapping( + value: object, +) -> TypeIs[Mapping[object, object]]: # guard-ok: isinstance decides it, keys and values stay object + return isinstance(value, Mapping) + + +def json_value(value: object) -> JsonValue: + plain: Final[object] = dict(value) if _is_mapping(value) else value + try: + return _JSON_VALUE.validate_python(plain) + except ValidationError: + return _JSON_VALUE.validate_json(safe_dumps(plain)) + + +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, 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: + 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 + 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) -> Mapping[str, JsonValue] | None: + try: + return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + + +def _stored_json_object(value: object) -> Mapping[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 decrypt_stored_json_object(value: object) -> Mapping[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 +) -> Mapping[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) -> Mapping[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/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 24b30e65e2b..7a967981cac 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -759,7 +759,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): if reject_invalid_logging_only_scope: @@ -1098,7 +1097,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 14dd9ce6e8a..97a3c93f670 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -56,6 +56,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( # noqa: F401 # decrypt_value_helper, encrypt_value_helper, get_salt_key, + parse_stored_json_object, ) ValueClass = Literal["migrated", "legacy", "plaintext", "undecryptable", "not-a-string"] @@ -470,6 +471,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",)), @@ -502,7 +504,9 @@ def _classify_into_report(report: LocationReport, value: str) -> None: names, base URLs, …) do not decrypt and fall through to ``plaintext``, so over-scanning a column is harmless to the residual count. """ - report.count(classify_value(value, key="scan")) + from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX + + report.count(classify_value(value.removeprefix(CALLBACK_VAR_ENCRYPTED_PREFIX), key="scan")) def _classify_secret_map(value: object, key: str) -> ValueClass | None: @@ -554,23 +558,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 = _DECODED_JSON.validate_python(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 @@ -580,7 +584,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 @@ -595,7 +600,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/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index 477cf44073f..7ff3f1d271b 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -22,6 +22,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import fallback_target_model_name 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 decrypt_config_section, encrypt_config_section from litellm.proxy.utils import PrismaClient, invalidate_config_param if TYPE_CHECKING: @@ -142,7 +143,7 @@ async def _stored_router_settings(prisma_client: PrismaClient) -> Mapping[str, J return {} if row is None or not isinstance(row.param_value, Mapping): return {} - return _STORED_ROUTER_SETTINGS.validate_python(row.param_value) + return decrypt_config_section(ROUTER_SETTINGS_PARAM, _STORED_ROUTER_SETTINGS.validate_python(row.param_value)) def _editable_rules( @@ -174,7 +175,7 @@ async def _write_rules( async def _persist_router_settings(prisma_client: PrismaClient, router_settings: Mapping[str, object]) -> None: - router_settings_json: Final = json.dumps(router_settings) + router_settings_json: Final = json.dumps(encrypt_config_section(ROUTER_SETTINGS_PARAM, router_settings)) await ConfigRepository(prisma_client).table.upsert( where={"param_name": ROUTER_SETTINGS_PARAM}, data={ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d28401524ed..043f1240c63 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -76,6 +76,11 @@ 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, + parse_stored_json_object, +) from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for 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 @@ -5917,6 +5922,27 @@ async def delete_key_aliases( ) +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 + 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(rekeyed)} + ) + + async def rotate_master_key( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, @@ -6001,6 +6027,10 @@ async def rotate_master_key( data={"param_value": prisma.Json({**environment_variables_dict, **encrypted_env_vars})}, ) + for c in config: + if c.param_name in ENCRYPTED_CONFIG_SECTIONS and c.param_value is not None: + await _rekey_config_section(prisma_client, c.param_name, c.param_value, new_master_key) + try: from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 2c7453ebb1c..4303b2dac4d 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -74,8 +74,9 @@ from litellm.proxy.common_utils.credential_hydration import ( write_touches_federation_surface, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_json_strings, decrypt_value_helper, - encrypt_value_helper, + json_value, ) from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -126,7 +127,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import ( from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.chunked_in import find_many_in 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.project_repository import ProjectRepository from litellm.repositories.table_repositories import ModelTableRepository @@ -510,8 +511,9 @@ def _effective_complexity_router_config( incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config existing: Final = None if existing_params is None else existing_params.complexity_router_config config_adapter: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + stored: Final = decrypt_json_strings(json_value(config_adapter.validate_python(existing))) return resolve_complexity_router_config_write( - config_adapter.validate_python(incoming), config_adapter.validate_python(existing) + config_adapter.validate_python(incoming), config_adapter.validate_python(stored) ).effective @@ -1432,19 +1434,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, exclude_unset=True).items() } - - merged_litellm_params.update(encrypted_params) + merged_litellm_params.update(encrypt_model_litellm_params({**merged_litellm_params, **patched_params})) # update model info if updated_patch.model_info: @@ -2062,11 +2061,8 @@ async def add_model_to_db( model_params.litellm_params, None ) _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) if isinstance(v, str) else v - ) - model_params.litellm_params[k] = encrypted_value + 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, "model_name": model_params.model_name, @@ -2557,6 +2553,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 @@ -2624,7 +2629,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 or () ): @@ -3369,25 +3374,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 = encrypt_model_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 + } + ) base_update: Final[PrismaCompatibleUpdateDBModel] = { "litellm_params": json.dumps(merged_dictionary), diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 6ace23d016b..0fa947809fb 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -22,6 +22,7 @@ from pydantic import TypeAdapter 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, @@ -32,7 +33,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity from litellm.repositories.daily_activity_buckets import api_key_from_bucket -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, @@ -365,7 +366,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: @@ -374,24 +374,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/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/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f837e75039e..3163a3178f4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -531,8 +531,13 @@ from litellm.proxy.common_utils.discoverable_model_filter import ( 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.fips import enforce_fips_mode @@ -6247,11 +6252,7 @@ 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) - else MappingProxyType({}) + MappingProxyType({}) if existing_value is None else decrypt_config_section(section_name, existing_value) ) merged_section: Final[Mapping[str, JsonValue]] = MappingProxyType( { @@ -6264,7 +6265,7 @@ class ProxyConfig: ) if merged_section == existing_section: return None - serialized_section: Final = json.dumps(dict(merged_section)) + 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}, @@ -7972,7 +7973,7 @@ class ProxyConfig: def _resolve_db_litellm_param(self, key: str, value: object) -> object: if not isinstance(value, str): - return value + return decrypt_json_strings(json_value(value)) decrypted_value: Final = decrypt_value_helper(value=value, key=key, return_original_value=True) if ( @@ -8367,7 +8368,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 ) @@ -8835,7 +8836,8 @@ class ProxyConfig: self._warned_unreadable_db_general_settings = True def _prepared_db_settings_values(self, section: Section, value: object) -> Mapping[str, SettingsJsonValue]: - 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_db_environment_variables( @@ -20454,10 +20456,10 @@ async def update_config( ) if row is None or row.param_value is None: return {} - return dict(row.param_value) + return dict(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, json.dumps(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 f81d3d2782f..d2ea54d18db 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -6,7 +6,7 @@ import json from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, JsonValue, TypeAdapter from litellm.models.model import LiteLLM_ProxyModelTable from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict @@ -23,6 +23,55 @@ class _ProxyModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_Pr table_name = "litellm_proxymodeltable" +CLASSIFIER_CREDENTIAL_FIELDS: Final = frozenset({"api_key", "api_base"}) +CLASSIFIER_CONFIG_KEYS: Final = frozenset({"opensource_classifier_config", "jev_classifier_config"}) + + +def _classifier_with_encrypted_credentials(classifier: JsonValue, new_encryption_key: str | None) -> JsonValue: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_json_strings + + if not isinstance(classifier, dict): + return classifier + return { + key: ( + encrypt_json_strings(value, new_encryption_key=new_encryption_key) + if key in CLASSIFIER_CREDENTIAL_FIELDS and isinstance(value, str) + else value + ) + for key, value in classifier.items() + } + + +def _complexity_router_config_with_encrypted_classifier_credentials( + config: JsonValue, new_encryption_key: str | None +) -> JsonValue: + if not isinstance(config, dict): + return config + return { + key: ( + _classifier_with_encrypted_credentials(value, new_encryption_key) + if key in CLASSIFIER_CONFIG_KEYS + else value + ) + for key, value in config.items() + } + + +def encrypt_model_litellm_params( + litellm_params: Mapping[str, object], new_encryption_key: str | None = None +) -> Mapping[str, JsonValue]: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_json_strings, json_value + + return { + 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) + ) + for key, value in litellm_params.items() + } + + class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): """Repository for proxy model database operations with encryption support.""" @@ -39,30 +88,15 @@ 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.""" - from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper - - 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 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]: - """Decrypt sensitive values in litellm_params.""" - from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, json_value - 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), signing_key=self._encryption_key) + for key, value in litellm_params.items() + } def _to_model(self, record: DbRecord | None) -> LiteLLM_ProxyModelTable | None: """Convert a database record to a Model with decryption.""" diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index f4ca7fa2965..384b253fd80 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -85,6 +85,8 @@ IGNORE_FUNCTIONS = [ "_redacted_client_field", # bounded by the nesting depth of a JSON-decoded error body (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. "_without_requester_ip_address", # walks only the fresh tree _sanitize_request_body_for_spend_logs_payload returns (mapping cycles already cut), so it never goes deeper than that sanitizer did. "_merge_session_fields", # max depth set (_SESSION_MERGE_MAX_DEPTH); deeper mappings are replaced wholesale by the later session.update. "_json", # router_rust index views: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising ValueError at the cap. diff --git a/tests/unit/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/unit/proxy/common_utils/test_encrypt_decrypt_utils.py index 1163090c510..1a503906394 100644 --- a/tests/unit/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/unit/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -7,21 +7,36 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130 """ import base64 +import json import re +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 from litellm.proxy.common_utils.encrypt_decrypt_utils import ( V2_GCM_PREFIX, decrypt_bearer_token, + decrypt_config_section, decrypt_if_encrypted_with, + decrypt_json_strings, + decrypt_stored_json_object, decrypt_value_helper, encrypt_bearer_token, + encrypt_config_section, + encrypt_json_strings, encrypt_value, encrypt_value_helper, + json_value, + parse_stored_json_object, ) +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.""" @@ -31,7 +46,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 @@ -241,6 +256,183 @@ def test_explicit_key_decrypt_supports_the_empty_master_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, +} + + +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]) +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 tuple(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_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): + deepest = deepest["child"] + assert deepest != "deep-secret" + assert decrypt_json_strings(stored) == value + + +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" + + +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_config_section_rekeys_a_row_stored_as_a_json_string(): + stored_row = json.dumps(encrypt_json_strings({"redis_password": "redis-secret", "num_retries": 2})) + + rotated = encrypt_config_section("router_settings", stored_row, new_encryption_key="sk-rotated") + + assert decrypt_if_encrypted_with(rotated["redis_password"], "sk-rotated") == "redis-secret" + assert rotated["num_retries"] == 2 + + +@pytest.mark.parametrize("stored_row", ("not json", json.dumps(["api_key"]), 7, None)) +def test_parse_stored_json_object_returns_none_for_a_row_that_is_not_a_json_object(stored_row): + assert parse_stored_json_object(stored_row) 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 + def test_bearer_token_opens_only_under_its_own_prefix(): token = encrypt_bearer_token("session", prefix="kind_a_") relabeled = "kind_b_" + token.removeprefix("kind_a_") diff --git a/tests/unit/proxy/db/test_master_key_migration.py b/tests/unit/proxy/db/test_master_key_migration.py index ebd76084e8d..fe41ce2fad7 100644 --- a/tests/unit/proxy/db/test_master_key_migration.py +++ b/tests/unit/proxy/db/test_master_key_migration.py @@ -99,6 +99,7 @@ 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)}}, @@ -141,7 +142,7 @@ def _seeded_tables() -> Tables: } -_VALUES_UNDER_THE_PREVIOUS_KEY = 8 +_VALUES_UNDER_THE_PREVIOUS_KEY = 9 @pytest.mark.asyncio @@ -160,6 +161,8 @@ 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" 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" diff --git a/tests/unit/proxy/guardrails/test_guardrail_registry.py b/tests/unit/proxy/guardrails/test_guardrail_registry.py index 7eabc24e090..01494008c24 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_registry.py +++ b/tests/unit/proxy/guardrails/test_guardrail_registry.py @@ -1982,3 +1982,48 @@ async def test_rotate_guardrail_params_gives_up_on_a_row_that_keeps_changing(mon assert rows_updated == 0 assert prisma_client.db.litellm_guardrailstable.update_many.await_count == GUARDRAIL_ROTATION_ATTEMPTS assert prisma_client.db.litellm_guardrailstable.find_unique.await_count == GUARDRAIL_ROTATION_ATTEMPTS - 1 + + +@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/unit/proxy/management_endpoints/test_credential_migration.py b/tests/unit/proxy/management_endpoints/test_credential_migration.py index ffb855173fa..eded5c0f586 100644 --- a/tests/unit/proxy/management_endpoints/test_credential_migration.py +++ b/tests/unit/proxy/management_endpoints/test_credential_migration.py @@ -16,6 +16,7 @@ import pytest from litellm._service_logger import ServiceTypes from litellm.proxy import proxy_server +from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX from litellm.proxy.common_utils.encrypt_decrypt_utils import ( V2_GCM_PREFIX, encrypt_value_helper, @@ -483,6 +484,27 @@ async def test_scan_covered_tables_classifies_search_tool_params(salt_key, monke assert by_loc["search_tools"].plaintext == 1 +@pytest.mark.asyncio +async def test_scan_covered_tables_classifies_marked_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": CALLBACK_VAR_ENCRYPTED_PREFIX + 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 == 1 + 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")) @@ -601,6 +623,24 @@ async def test_migrate_covered_tables_reports_real_counts(salt_key, monkeypatch) 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 + def _db_service_hooks() -> tuple[AsyncMock, MagicMock]: success: Final = AsyncMock() service_logging: Final = MagicMock(async_service_success_hook=success, async_service_failure_hook=AsyncMock()) 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 b453f985512..431e72a3cd9 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -9811,6 +9811,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( @@ -19116,6 +19117,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=[] ) @@ -23570,6 +23572,139 @@ class TestTeamAdminMemberKeyBudgetUpdate: assert "member_key_budgets" not in str(exc.value.detail) +def _rotation_prisma_client(stored_router_settings: object) -> AsyncMock: + 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=[]) + mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.writer_db = mock_prisma_client.db + 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, + user_api_key_dict=UserAPIKeyAuth( + 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", + ) + + +@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(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}) + mock_prisma_client = _rotation_prisma_client(stored_router_settings) + 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) + 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 + + +@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_stored_as_a_json_string(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"})) + mock_prisma_client = _rotation_prisma_client(stored_router_settings) + 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" + + +@pytest.mark.asyncio +@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key") +async def test_rotate_master_key_reports_router_settings_that_are_not_a_json_object_without_its_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") + mock_rotate_mcp.return_value = None + + await _rotate_to_new_master_key(mock_prisma_client, monkeypatch) + + skipped: Final = [ + 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"] + assert "plaintext-router-secret" not in caplog.text + mock_prisma_client.db.litellm_config.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_router_settings_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) + mock_rotate_mcp.return_value = None + + await _rotate_to_new_master_key(mock_prisma_client, monkeypatch) + + skipped: Final = [ + 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}"] + assert "vendor-secret" not in caplog.text + mock_prisma_client.db.litellm_config.update.assert_not_called() + + _MEMBER_KEY_TEAM = LiteLLM_TeamTableCachedObj( team_id="team-1", members_with_roles=[ 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 32c9535b4c8..e67c2e36fb2 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -30,7 +30,11 @@ 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_json_strings, + encrypt_value_helper, + json_value, +) from litellm.proxy.management_endpoints.model_management_endpoints import ( KubernetesPodsResponse, ModelManagementAuthChecks, @@ -819,6 +823,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 @@ -2045,8 +2075,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", @@ -2096,8 +2126,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", @@ -3140,7 +3170,7 @@ class TestAddAndDeleteModelLifecycle: mock_router.registry.return_value.names.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), @@ -4801,7 +4831,7 @@ class TestModelInfoServerDerivedPricingFilter: mock_router.registry.return_value.names.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), @@ -5212,7 +5242,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"]) @@ -5241,7 +5271,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"]) @@ -5260,7 +5290,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"]) @@ -5288,7 +5318,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"]) @@ -5315,7 +5345,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"]) @@ -5391,7 +5421,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( @@ -7140,7 +7170,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, ), ): @@ -7206,7 +7236,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 @@ -7246,7 +7276,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, ), ): @@ -8622,7 +8652,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())) @@ -8990,7 +9020,7 @@ class TestTeamMemberAutoRouterWrites: await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor)) written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"] saved: Final = json.loads(written["litellm_params"])["complexity_router_config"] - assert saved == ( + assert decrypt_json_strings(json_value(saved)) == ( self._classifier_config({**connection, "timeout_ms": 8100}, False) if change_config else stored_config ) @@ -9079,7 +9109,8 @@ class TestTeamMemberAutoRouterWrites: await add_new_model(deployment, actor) written: Final = database.db.litellm_proxymodeltable.create.await_args.kwargs["data"] saved: Final = json.loads(written["litellm_params"])["complexity_router_config"] - assert saved == { + assert "stored-secret" not in written["litellm_params"] + assert decrypt_json_strings(json_value(saved)) == { "classifier_type": "oss_classifier", "tiers": {"SIMPLE": "allowed"}, "opensource_classifier_config": {**classifier, "provider": "laya" if provider == "laya" else "jev"}, @@ -9157,29 +9188,34 @@ class TestTeamMemberAutoRouterWrites: @pytest.mark.asyncio @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) @pytest.mark.parametrize("stored_legacy,supplied_legacy", [(True, True), (True, False), (False, True), (False, False)]) - @pytest.mark.parametrize("change", ["save", "rotate", "move", "move-without-key", "reset", "heuristic"]) + @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, change: str, stored_legacy: bool, supplied_legacy: bool + self, endpoint: str, change: str, stored_legacy: bool, supplied_legacy: bool, stored_as: 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 = self._classifier_config( - {**transport, "instructions": "Old instructions", "timeout_ms": 6100}, stored_legacy - ) - row: Final = original.model_copy( - update={ - "litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": stored_config, - }, - } + plaintext_params: Final = { + "model": "auto_router/complexity_router", + "complexity_router_config": self._classifier_config( + {**transport, "instructions": "Old instructions", "timeout_ms": 6100}, stored_legacy + ), + } + 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] @@ -9205,17 +9241,23 @@ class TestTeamMemberAutoRouterWrites: await operation written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"] saved: Final = json.loads(written["litellm_params"])["complexity_router_config"] + expected_classifier: Final = {**transport, "timeout_ms": 8100, **overrides} expected: Final = ( config if change == "heuristic" else { "classifier_type": "oss_classifier", "tiers": {"SIMPLE": "allowed"}, - "opensource_classifier_config": {**transport, "timeout_ms": 8100, **overrides}, + "opensource_classifier_config": expected_classifier, } ) - assert saved == expected - assert row.litellm_params["complexity_router_config"] == stored_config + stored_classifier: Final = saved.get("opensource_classifier_config") + if isinstance(stored_classifier, dict): + for field in ("api_key", "api_base"): + if isinstance(stored_classifier.get(field), str): + assert stored_classifier[field] != expected_classifier[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 @@ -9282,7 +9324,7 @@ class TestTeamMemberAutoRouterWrites: written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"] saved: Final = json.loads(written["litellm_params"])["complexity_router_config"] expected_provider: Final = supplied.get("provider", stored_provider) - assert saved == { + assert decrypt_json_strings(json_value(saved)) == { "classifier_type": "oss_classifier", "tiers": {"SIMPLE": "allowed"}, "opensource_classifier_config": { @@ -9361,7 +9403,7 @@ class TestTeamMemberAutoRouterWrites: return response: Final = 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_config: Final = json.loads(written["litellm_params"])["complexity_router_config"] + saved_config: Final = decrypt_json_strings(json_value(json.loads(written["litellm_params"])["complexity_router_config"])) untouched: Final = config_shape in ("omit-config", "null-config") saved: Final = saved_config["jev_classifier_config" if untouched else "opensource_classifier_config"] assert saved == secret_params["complexity_router_config"]["jev_classifier_config"] @@ -9385,6 +9427,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"]) @@ -9610,8 +9683,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", @@ -9659,8 +9732,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", @@ -9781,8 +9854,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", @@ -9840,6 +9913,122 @@ class TestModelManagementActorEdges: 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 = AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)) + mock_router: Final = MagicMock() + mock_router.registry.return_value.names.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 + class TestAddModelToDbBlocked: """`_add_model_to_db` must thread `blocked` into the initial insert, so the wizard can create a discovered-but-unchecked model already paused instead of active-then-patched.""" @@ -10300,8 +10489,8 @@ class TestNonAdminCannotPersistWifFieldsOnModel: True, ), patch( # test-quality-ok: the proxy wiring under test is what this patches - "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, **_: value, ), patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -10441,8 +10630,8 @@ class TestNonAdminCannotPersistWifFieldsOnModel: patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test patch( # test-quality-ok: the proxy wiring under test is what this patches - "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, **_: value, ), patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -10476,8 +10665,8 @@ class TestNonAdminCannotPersistWifFieldsOnModel: patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test patch( # test-quality-ok: the proxy wiring under test is what this patches - "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, **_: value, ), patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -10649,8 +10838,8 @@ class TestNonAdminCannotPersistWifFieldsOnModel: True, ), patch( # test-quality-ok: the proxy wiring under test is what this patches - "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, **_: value, ), patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -10931,8 +11120,8 @@ class TestWifBoundaryReadsTheResultingDeployment: patch("litellm.proxy.proxy_server.store_model_in_db", True), patch("litellm.proxy.proxy_server.premium_user", True), 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, **_: value, ), patch( "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", @@ -11073,8 +11262,8 @@ class TestFederationGateScopesToWhatTheWriteTouches: patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test patch( # test-quality-ok: proxy wiring under test - "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, **_: value, ), patch( # test-quality-ok: proxy wiring under test "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index 531bf98b60e..a7ad8121812 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -1498,58 +1498,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(): """ @@ -1585,6 +1533,40 @@ async def test_add_tag_to_deployment_model_not_found(): 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"]} + class _StoredTagTable: def __init__(self, model_info: object) -> None: self.model_info = model_info 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 diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 712b4134c05..9cb856a8c7f 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -1345,6 +1345,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/unit/proxy/test_fallback_management_endpoints.py b/tests/unit/proxy/test_fallback_management_endpoints.py index e3d0d072e0a..0896fcfaa42 100644 --- a/tests/unit/proxy/test_fallback_management_endpoints.py +++ b/tests/unit/proxy/test_fallback_management_endpoints.py @@ -22,6 +22,7 @@ from redis.exceptions import NoPermissionError import litellm from litellm import Router +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, encrypt_json_strings from litellm.proxy.config_resolvers.settings_rules import Section from litellm.proxy.management_endpoints.fallback_management_endpoints import ( FallbackCreateRequest, @@ -344,9 +345,11 @@ class _RedisThatRefusesDel: class TestFallbackWritesSeeTheLatestStoredRules: """A write reads the rules the database holds now and leaves no stale copy in the config cache behind""" - async def _create(self, request: FallbackCreateRequest, stored: _StoredRouterSettings) -> FallbackResponse: + async def _create( + self, request: FallbackCreateRequest, stored: _StoredRouterSettings, router: Router | None = None + ) -> FallbackResponse: with ( - patch("litellm.proxy.proxy_server.llm_router", _team_router()), + patch("litellm.proxy.proxy_server.llm_router", router or _team_router()), patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config), patch("litellm.proxy.proxy_server.store_model_in_db", True), @@ -369,6 +372,28 @@ class TestFallbackWritesSeeTheLatestStoredRules: assert stored.stored_fallbacks() == [TEAM_RULE, GATEWAY_RULE] + async def test_a_create_beside_an_encrypted_row_applies_and_stores_its_rules_decrypted( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-fallback-salt-1234") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + stored: Final = _StoredRouterSettings( + encrypt_json_strings({"redis_password": "redis-canary-password", "fallbacks": [TEAM_RULE]}) + ) + router: Final = _team_router() + + await self._create( + FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored, router + ) + + assert stored.row is not None + assert "redis-canary-password" not in json.dumps(stored.row.param_value) + assert decrypt_json_strings(stored.row.param_value) == { + "redis_password": "redis-canary-password", + "fallbacks": [TEAM_RULE, GATEWAY_RULE], + } + assert router.settings().fallbacks == [TEAM_RULE, GATEWAY_RULE] + async def test_a_second_create_keeps_the_first_rule(self) -> None: stored: Final = _StoredRouterSettings(None) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 8536cbaa399..c0411baaab7 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -16538,6 +16538,59 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is ] +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 = route_entry_points_to_stubs(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" + + def test_feedback_box_is_one_json_record_in_json_mode( monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] ) -> None: diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py index 56843640888..ab77dc1db71 100644 --- a/tests/unit/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -431,11 +431,8 @@ class TestModelRepository: "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: f"encrypted_{v}", ) - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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"}, @@ -449,11 +446,7 @@ class TestModelRepository: "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: f"encrypted_{v}", ) - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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={ @@ -473,11 +466,7 @@ class TestModelRepository: "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: f"encrypted_{v}", ) - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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", @@ -495,11 +484,7 @@ class TestModelRepository: assert updated.model_name == "new-name" @pytest.mark.asyncio - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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", @@ -518,11 +503,7 @@ class TestModelRepository: assert len(models) == 2 @pytest.mark.asyncio - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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", @@ -535,11 +516,7 @@ class TestModelRepository: assert len(models) == 1 @pytest.mark.asyncio - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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", @@ -555,11 +532,7 @@ class TestModelRepository: "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: v, ) - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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", @@ -574,11 +547,7 @@ class TestModelRepository: assert updated.blocked is True @pytest.mark.asyncio - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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", @@ -592,11 +561,7 @@ class TestModelRepository: "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda v, **kw: v, ) - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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", @@ -2056,11 +2021,7 @@ class TestModelRepositoryExtended: return ModelRepository(client) @pytest.mark.asyncio - @patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils.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", @@ -2485,6 +2446,91 @@ class TestAutoRouterSessionRepository: _ = 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 + + @pytest.mark.asyncio + 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", + "instructions": "Rate the question", + }, + } + 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 + 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"]["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 + + @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" + + class TestSubscriptionAccountRepository: @pytest.mark.asyncio async def test_update_fee_writes_only_the_named_columns(self):