diff --git a/.circleci/config.yml b/.circleci/config.yml index 602604714bd..4d38d0abe27 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1650,6 +1650,7 @@ jobs: command: | docker run -d \ -p 4001:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL="postgresql://postgres:postgres@host.docker.internal:5432/litellm_test" \ -e LITELLM_MASTER_KEY="sk-1234" \ --name schema-seed \ @@ -1670,6 +1671,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL="postgresql://postgres:postgres@host.docker.internal:5432/litellm_test" \ -e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \ -e DISABLE_SCHEMA_UPDATE="True" \ @@ -1744,6 +1746,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e USE_PRISMA_MIGRATE=True \ -e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \ @@ -1839,6 +1842,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e AZURE_API_KEY=$AZURE_API_KEY \ -e AZURE_API_BASE=$AZURE_API_BASE \ @@ -1927,6 +1931,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e REDIS_HOST=$REDIS_HOST \ -e REDIS_PASSWORD=$REDIS_PASSWORD \ @@ -1987,6 +1992,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e REDIS_HOST=$REDIS_HOST \ -e REDIS_PASSWORD=$REDIS_PASSWORD \ @@ -2064,6 +2070,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e REDIS_HOST=host.docker.internal \ -e REDIS_PORT=6379 \ @@ -2146,6 +2153,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e REDIS_HOST=$REDIS_HOST \ -e REDIS_PASSWORD=$REDIS_PASSWORD \ @@ -2168,6 +2176,7 @@ jobs: command: | docker run -d \ -p 4001:4001 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e REDIS_HOST=$REDIS_HOST \ -e REDIS_PASSWORD=$REDIS_PASSWORD \ @@ -2245,6 +2254,7 @@ jobs: docker run -d \ --restart on-failure \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e STORE_MODEL_IN_DB="True" \ -e LITELLM_MASTER_KEY="sk-1234" \ @@ -2319,6 +2329,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e REDIS_HOST=$REDIS_HOST \ -e REDIS_PASSWORD=$REDIS_PASSWORD \ @@ -2401,6 +2412,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e LITELLM_MASTER_KEY="sk-1234" \ -e OPENAI_API_KEY=$OPENAI_API_KEY \ @@ -2492,6 +2504,7 @@ jobs: command: | docker run -d \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \ -e LITELLM_MASTER_KEY="sk-1234" \ -e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \ @@ -2673,6 +2686,7 @@ jobs: name: Start LiteLLM proxy environment: LITELLM_MASTER_KEY: "sk-1234" + LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY: "true" MOCK_LLM_URL: "http://127.0.0.1:8090/v1" DISABLE_SCHEMA_UPDATE: "true" SERVER_ROOT_PATH: "" @@ -2816,6 +2830,7 @@ jobs: name: Start LiteLLM proxy under a server root path environment: LITELLM_MASTER_KEY: "sk-1234" + LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY: "true" MOCK_LLM_URL: "http://127.0.0.1:8090/v1" DISABLE_SCHEMA_UPDATE: "true" # Output flows to this step's own log, so a boot crash is visible here @@ -2901,6 +2916,7 @@ jobs: command: | docker run --name my-app \ -p 4000:4000 \ + -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \ -e DATABASE_URL="postgresql://wrong:wrong@wrong:5432/wrong" \ myapp:latest \ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d24e451cfc1..9344b982adc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2595,6 +2595,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): use_google_kms: bool | None = Field(None, description="decrypt keys with google kms") use_azure_key_vault: bool | None = Field(None, description="load keys from azure key vault") master_key: str | None = Field(None, description="require a key for all calls to proxy") + dangerously_permit_weak_or_unset_master_key: bool | None = Field( + None, + description="local development only: start even when master_key is unset, empty, or a publicly known default", + ) coordination_redis: CoordinationRedisParams | None = Field( None, description=( diff --git a/litellm/proxy/auth/master_key_boot_check.py b/litellm/proxy/auth/master_key_boot_check.py new file mode 100644 index 00000000000..aa50c4e34c1 --- /dev/null +++ b/litellm/proxy/auth/master_key_boot_check.py @@ -0,0 +1,291 @@ +import atexit +import sys +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, replace +from enum import Enum +from types import MappingProxyType +from typing import Final + +from typing_extensions import assert_never + +from litellm._logging import verbose_proxy_logger + +WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING: Final = "dangerously_permit_weak_or_unset_master_key" +WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR: Final = "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY" +MASTER_KEY_SETTING: Final = "master_key" +MASTER_KEY_ENV_VAR: Final = "LITELLM_MASTER_KEY" +SALT_KEY_ENV_VAR: Final = "LITELLM_SALT_KEY" +MIGRATE_FROM_MASTER_KEY_ENV_VAR: Final = "LITELLM_MIGRATE_FROM_MASTER_KEY" +PUBLICLY_KNOWN_MASTER_KEYS: Final = frozenset({"sk-1234"}) +ROTATION_DOCS_URL: Final = "https://docs.litellm.ai/docs/proxy/master_key_rotations#proxy-refuses-to-start" +_NEW_MASTER_KEY: Final = "sk-$(openssl rand -hex 32)" +GENERATE_MASTER_KEY_COMMAND: Final = f'echo "{MASTER_KEY_ENV_VAR}={_NEW_MASTER_KEY}" | tee -a .env' +PRINT_NEW_MASTER_KEY_COMMAND: Final = f'echo "{_NEW_MASTER_KEY}"' + + +class UnsafeMasterKeyReason(Enum): + NOT_SET = "not_set" + EMPTY = "empty" + PUBLICLY_KNOWN = "publicly_known" + + +@dataclass(frozen=True, slots=True) +class ConfigFileSource: + config_file_path: str | None + + +@dataclass(frozen=True, slots=True) +class EnvironmentSource: + pass + + +MasterKeySource = ConfigFileSource | EnvironmentSource + + +@dataclass(frozen=True, slots=True) +class SafeMasterKey: + pass + + +@dataclass(frozen=True, slots=True) +class UnsafeMasterKeyAllowed: + reason: UnsafeMasterKeyReason + + +@dataclass(frozen=True, slots=True) +class StoredSecretsMigration: + from_master_key: str + encrypted_value_count: int | None + + +@dataclass(frozen=True, slots=True) +class UnsafeMasterKeyRefused: + reason: UnsafeMasterKeyReason + source: MasterKeySource + environment_variable_is_set: bool + migration: StoredSecretsMigration | None + + +MasterKeyBootVerdict = SafeMasterKey | UnsafeMasterKeyAllowed | UnsafeMasterKeyRefused + + +class UnsafeMasterKeyError(Exception): + pass + + +def master_key_boot_verdict( + *, + master_key: str | None, + environment_master_key: str | None, + general_settings: Mapping[str, object], + config_file_path: str | None, + override_env_is_on: bool, + salt_key_is_set: bool, + database_is_configured: bool, +) -> MasterKeyBootVerdict: + reason: Final = _unsafe_reason(master_key) + if reason is None: + return SafeMasterKey() + if override_env_is_on or general_settings.get(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING) is True: + return UnsafeMasterKeyAllowed(reason=reason) + config_file_only_relays_the_environment: Final = master_key is not None and master_key == environment_master_key + return UnsafeMasterKeyRefused( + reason=reason, + source=( + ConfigFileSource(config_file_path=config_file_path) + if MASTER_KEY_SETTING in general_settings and not config_file_only_relays_the_environment + else EnvironmentSource() + ), + environment_variable_is_set=environment_master_key is not None, + migration=( + StoredSecretsMigration(from_master_key=master_key, encrypted_value_count=None) + if master_key is not None and not salt_key_is_set and database_is_configured + else None + ), + ) + + +async def with_stored_secrets_counted( + verdict: MasterKeyBootVerdict, count_values_encrypted_with: Callable[[str], Awaitable[int | None]] +) -> MasterKeyBootVerdict: + if not isinstance(verdict, UnsafeMasterKeyRefused) or verdict.migration is None: + return verdict + count: Final = await count_values_encrypted_with(verdict.migration.from_master_key) + return replace(verdict, migration=None if count == 0 else replace(verdict.migration, encrypted_value_count=count)) + + +def enforce_master_key_boot_verdict(verdict: MasterKeyBootVerdict, announce: Callable[[str], object]) -> None: + match verdict: + case SafeMasterKey(): + return + case UnsafeMasterKeyAllowed(reason=reason): + verbose_proxy_logger.warning( + "%s is on, so the proxy is starting with %s. Never run this outside local development.", + WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING, + _UNSAFE_STATE[reason], + ) + case UnsafeMasterKeyRefused(reason=reason): + announce(f"\n{render_refusal(verdict)}\n\n") + raise UnsafeMasterKeyError( + f"LiteLLM proxy refused to start: {_REFUSAL_HEADLINE[reason]} The fix is printed once the server exits." + ) + case _: + assert_never(verdict) + + +def announce_on_stderr_at_exit(message: str) -> None: + """A logger would redact the key-shaped command and the lifespan traceback would bury it, so print at exit.""" + atexit.register(_flush_stdout_then_write_stderr, message) + + +def _flush_stdout_then_write_stderr(message: str) -> None: + sys.stdout.flush() + sys.stderr.write(message) + + +def render_refusal(refusal: UnsafeMasterKeyRefused) -> str: + return "\n\n".join( + ( + f"LiteLLM proxy refused to start: {_REFUSAL_HEADLINE[refusal.reason]}\n{_source_line(refusal)}", + _fix_steps(refusal), + _OVERRIDE_HINT, + ) + ) + + +_UNSAFE_STATE: Final = MappingProxyType( + { + UnsafeMasterKeyReason.NOT_SET: "no master key, which accepts every request without authentication", + UnsafeMasterKeyReason.EMPTY: "an empty master key", + UnsafeMasterKeyReason.PUBLICLY_KNOWN: "a publicly known master key", + } +) + +_REFUSAL_HEADLINE: Final = MappingProxyType( + { + UnsafeMasterKeyReason.NOT_SET: ( + "no master key is set, so every request would be accepted without authentication." + ), + UnsafeMasterKeyReason.EMPTY: "the master key is empty.", + UnsafeMasterKeyReason.PUBLICLY_KNOWN: "the master key is a publicly known default.", + } +) + +_SAVE_KEY_STEP: Final = ( + "Generate a key and save it to .env:\n" + f" {GENERATE_MASTER_KEY_COMMAND}\n" + " Not using a .env file (docker run, Kubernetes, pip install)? Pass the same value as the\n" + f" {MASTER_KEY_ENV_VAR} environment variable instead." +) + +_REPLACE_EXPORTED_KEY_STEP: Final = ( + "Generate a key:\n" + f" {PRINT_NEW_MASTER_KEY_COMMAND}\n" + f" Put it in place of the current {MASTER_KEY_ENV_VAR} value wherever that is set: a shell export, your\n" + " container or deployment environment, or its line in .env. Do not just add it to .env, because a value\n" + " already exported in the environment wins over .env." +) + +_RESTART_TO_MIGRATE_STEP: Final = ( + "Start the proxy again. It re-encrypts the stored values with the new key, then logs that\n" + f" {MIGRATE_FROM_MASTER_KEY_ENV_VAR} can be removed. Details: {ROTATION_DOCS_URL}" +) + +_OVERRIDE_HINT: Final = ( + f"Local development only: set {WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR}=true, or\n" + f"general_settings.{WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING}: true, to start anyway." +) + + +def _unsafe_reason(master_key: str | None) -> UnsafeMasterKeyReason | None: + if master_key is None: + return UnsafeMasterKeyReason.NOT_SET + stripped: Final = master_key.strip() + if not stripped: + return UnsafeMasterKeyReason.EMPTY + if stripped in PUBLICLY_KNOWN_MASTER_KEYS: + return UnsafeMasterKeyReason.PUBLICLY_KNOWN + return None + + +def _config_label(source: ConfigFileSource) -> str: + return source.config_file_path or "your config" + + +def _source_line(refusal: UnsafeMasterKeyRefused) -> str: + match refusal.source: + case ConfigFileSource() as source: + if refusal.reason is UnsafeMasterKeyReason.NOT_SET: + return ( + f"general_settings.{MASTER_KEY_SETTING} in {_config_label(source)} is blank, or points at an " + "environment variable that is not set." + ) + return f"It comes from general_settings.{MASTER_KEY_SETTING} in {_config_label(source)}." + case EnvironmentSource(): + if refusal.reason is UnsafeMasterKeyReason.NOT_SET: + return ( + f"Neither general_settings.{MASTER_KEY_SETTING} nor the {MASTER_KEY_ENV_VAR} " + "environment variable is set." + ) + return f"It comes from the {MASTER_KEY_ENV_VAR} environment variable." + case _: + assert_never(refusal.source) + + +def _fix_steps(refusal: UnsafeMasterKeyRefused) -> str: + steps: Final = (*_config_steps(refusal.source), *_key_steps(refusal)) + numbered: Final = "\n".join(f"{number}. {step}" for number, step in enumerate(steps, start=1)) + return numbered if refusal.migration is None else f"{_migration_lead(refusal.migration)}\n{numbered}" + + +def _config_steps(source: MasterKeySource) -> tuple[str, ...]: + match source: + case ConfigFileSource(): + return ( + f"Make sure {_config_label(source)} reads the key from the environment:\n" + f" general_settings:\n {MASTER_KEY_SETTING}: os.environ/{MASTER_KEY_ENV_VAR}", + ) + case EnvironmentSource(): + return () + case _: + assert_never(source) + + +def _key_steps(refusal: UnsafeMasterKeyRefused) -> tuple[str, ...]: + if refusal.migration is None: + return (_REPLACE_EXPORTED_KEY_STEP if refusal.environment_variable_is_set else _SAVE_KEY_STEP,) + if refusal.environment_variable_is_set: + return ( + f"Set the key to migrate from next to {MASTER_KEY_ENV_VAR}, wherever that is set (a shell export, your\n" + " container or deployment environment, or .env):\n" + f" {_migrate_from_assignment(refusal.migration)}", + _REPLACE_EXPORTED_KEY_STEP, + _RESTART_TO_MIGRATE_STEP, + ) + return ( + "Save the key to migrate from and a newly generated key to .env:\n" + f" echo '{_migrate_from_assignment(refusal.migration)}' | tee -a .env\n" + f" {GENERATE_MASTER_KEY_COMMAND}\n" + " Not using a .env file (docker run, Kubernetes, pip install)? Pass the same two values as\n" + " environment variables instead.", + _RESTART_TO_MIGRATE_STEP, + ) + + +def _migrate_from_assignment(migration: StoredSecretsMigration) -> str: + key: Final = migration.from_master_key + value: Final = key if key == key.strip() else f'"{key}"' + return f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}={value}" + + +def _migration_lead(migration: StoredSecretsMigration) -> str: + found: Final = ( + "could not be checked for values" + if migration.encrypted_value_count is None + else f"holds {migration.encrypted_value_count} value(s)" + ) + return ( + f"Your database {found} encrypted with this master key,\n" + f"which encrypts stored credentials while {SALT_KEY_ENV_VAR} is not set. Replacing the key alone makes them\n" + "unreadable, so also tell the proxy which key to migrate from:" + ) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 561a53409f4..bdf45ad46f8 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -44,7 +44,8 @@ _EXTRA_SENSITIVE_CALLBACK_KEYS: Final = {"gcs_path_service_account"} # Sentinel prefix on encrypted callback_var values. Lets us detect # already-encrypted input cheaply (no decrypt-attempt round trip) and # avoid double-encrypting if `LITELLM_SALT_KEY` is rotated between writes. -_CALLBACK_VAR_ENCRYPTED_PREFIX: Final = "litellm_enc::" +CALLBACK_VAR_ENCRYPTED_PREFIX: Final = "litellm_enc::" +_CALLBACK_VAR_ENCRYPTED_PREFIX: Final = CALLBACK_VAR_ENCRYPTED_PREFIX # Metadata slots that hold operator-configured callback and secret-manager setup # (and therefore integration credentials). Resolved from UserAPIKeyAuth during # pre-call setup, never read back off the copies stamped into request metadata. diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 288dedebbc6..3584aaaf833 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -119,6 +119,35 @@ def encrypt_value_helper(value: str, new_encryption_key: str | None = None): raise e +def _legacy_ciphertext_bytes(value: str) -> bytes: + # Try URL-safe base64 decoding first (new format) + # Fall back to standard base64 decoding for backwards compatibility (old format) + try: + return base64.urlsafe_b64decode(value) + except Exception: + return base64.b64decode(value) + + +def _decrypt_with_signing_key(value: str, signing_key: str) -> str: + # Versioned AES-256-GCM values are detected before any base64 decode. + # The prefix is the algorithm tag the legacy nacl format never carried. + if value.startswith(_V2_GCM_PREFIX): + return _decrypt_aes_gcm(value=value, signing_key=signing_key) + + return decrypt_value(value=_legacy_ciphertext_bytes(value), signing_key=signing_key) + + +def decrypt_if_encrypted_with(value: str, signing_key: str) -> str | None: + """None unless value is a ciphertext under signing_key.""" + try: + # base64 decoding skips characters outside its alphabet, so "" and "*" decode to no bytes, + # which decrypt_value reads as an empty plaintext under any key. + decodes_to_nothing: Final = not value.startswith(_V2_GCM_PREFIX) and not _legacy_ciphertext_bytes(value) + return None if decodes_to_nothing else _decrypt_with_signing_key(value=value, signing_key=signing_key) + except Exception: # noqa: BLE001 # base64, nacl and AES-GCM each raise their own "not a ciphertext" type + return None + + def decrypt_value_helper( value: str, key: str, # this is just for debug purposes, showing the k,v pair that's invalid. not a signing key. @@ -129,21 +158,7 @@ def decrypt_value_helper( try: if isinstance(value, str): - # Versioned AES-256-GCM values are detected before any base64 decode. - # The prefix is the algorithm tag the legacy nacl format never carried. - if value.startswith(_V2_GCM_PREFIX): - return _decrypt_aes_gcm(value=value, signing_key=cast(str, signing_key)) - - # Try URL-safe base64 decoding first (new format) - # Fall back to standard base64 decoding for backwards compatibility (old format) - try: - decoded_b64 = base64.urlsafe_b64decode(value) - except Exception: - # If URL-safe decoding fails, try standard base64 decoding for backwards compatibility - decoded_b64 = base64.b64decode(value) - - value = decrypt_value(value=decoded_b64, signing_key=signing_key) - return value + return _decrypt_with_signing_key(value=value, signing_key=cast(str, signing_key)) # if it's not str - do not decrypt it, return the value return value diff --git a/litellm/proxy/db/master_key_migration.py b/litellm/proxy/db/master_key_migration.py new file mode 100644 index 00000000000..d100554201a --- /dev/null +++ b/litellm/proxy/db/master_key_migration.py @@ -0,0 +1,307 @@ +import json +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from enum import Enum +from typing import Final + +from pydantic import JsonValue, TypeAdapter +from typing_extensions import assert_never + +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.proxy.auth.master_key_boot_check import MIGRATE_FROM_MASTER_KEY_ENV_VAR, SALT_KEY_ENV_VAR +from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_value_helper +from litellm.proxy.db.create_views import SupportsRawQueries + + +@dataclass(frozen=True, slots=True) +class _SecretColumn: + table: str + primary_key: str + column: str + is_json: bool = True + only_rows_with_marked_ciphertexts: bool = False + + +_SECRET_COLUMNS: Final = ( + _SecretColumn("LiteLLM_ProxyModelTable", "model_id", "litellm_params"), + _SecretColumn("LiteLLM_CredentialsTable", "credential_id", "credential_values"), + _SecretColumn("LiteLLM_Config", "param_name", "param_value"), + _SecretColumn("LiteLLM_SSOConfig", "id", "sso_settings"), + _SecretColumn("LiteLLM_CacheConfig", "id", "cache_settings"), + _SecretColumn("LiteLLM_ConfigOverrides", "config_type", "config_value"), + _SecretColumn("LiteLLM_MCPServerTable", "server_id", "credentials"), + _SecretColumn("LiteLLM_MCPServerTable", "server_id", "static_headers"), + _SecretColumn("LiteLLM_MCPServerTable", "server_id", "env_vars"), + _SecretColumn("LiteLLM_MCPServerTable", "server_id", "env"), + _SecretColumn("LiteLLM_MCPServerOAuthClient", "server_id", "credentials"), + _SecretColumn("LiteLLM_MCPUserCredentials", "id", "credential_b64", is_json=False), + _SecretColumn("LiteLLM_MCPUserEnvVars", "id", "values_b64", is_json=False), + _SecretColumn("LiteLLM_SSOIdentityAssertion", "user_id", "assertion_b64", is_json=False), + _SecretColumn("LiteLLM_TeamTable", "team_id", "metadata", only_rows_with_marked_ciphertexts=True), + _SecretColumn("LiteLLM_VerificationToken", "token", "metadata", only_rows_with_marked_ciphertexts=True), + _SecretColumn("LiteLLM_UserTable", "user_id", "metadata", only_rows_with_marked_ciphertexts=True), + _SecretColumn("LiteLLM_DeletedTeamTable", "id", "metadata", only_rows_with_marked_ciphertexts=True), + _SecretColumn("LiteLLM_DeletedVerificationToken", "id", "metadata", only_rows_with_marked_ciphertexts=True), +) + +_STORED_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_PRIMARY_KEY: Final = TypeAdapter(str) + +ReplaceCiphertext = Callable[[str], str | None] + + +def replace_ciphertexts(value: JsonValue, replacement_for: ReplaceCiphertext, depth: int = 0) -> tuple[JsonValue, int]: + if depth > DEFAULT_MAX_RECURSE_DEPTH: + return value, 0 + match value: + case str(): + marker: Final = CALLBACK_VAR_ENCRYPTED_PREFIX if value.startswith(CALLBACK_VAR_ENCRYPTED_PREFIX) else "" + replacement: Final = replacement_for(value.removeprefix(marker)) + return (value, 0) if replacement is None else (marker + replacement, 1) + case list(): + items: Final = tuple(replace_ciphertexts(item, replacement_for, depth + 1) for item in value) + return [item for item, _ in items], sum(count for _, count in items) + case dict(): + fields: Final = {key: replace_ciphertexts(item, replacement_for, depth + 1) for key, item in value.items()} + return {key: item for key, (item, _) in fields.items()}, sum(count for _, count in fields.values()) + case _: + return value, 0 + + +async def count_values_encrypted_with(database: SupportsRawQueries, signing_key: str) -> int: + def keep(value: str) -> str | None: + return None if decrypt_if_encrypted_with(value, signing_key) is None else value + + return sum( + [ + replace_ciphertexts(_STORED_VALUE.validate_python(row[secret_column.column]), keep)[1] + for secret_column in await _secret_columns_in(database) + for row in await _rows_of(database, secret_column) + ] + ) + + +async def count_values_encrypted_with_or_none( + connect: Callable[[], Awaitable[SupportsRawQueries]], signing_key: str +) -> int | None: + try: + return await count_values_encrypted_with(await connect(), signing_key) + except Exception: # noqa: BLE001 # an unreadable database must not replace the boot refusal with a traceback + return None + + +async def reencrypt_stored_values(database: SupportsRawQueries, *, from_key: str, to_key: str) -> int: + def reencrypted(value: str) -> str | None: + plaintext: Final = decrypt_if_encrypted_with(value, from_key) + return None if plaintext is None else _CIPHERTEXT.validate_python(encrypt_value_helper(plaintext, to_key)) + + return sum( + [ + await _reencrypt_row(database, secret_column, row, reencrypted) + for secret_column in await _secret_columns_in(database) + for row in await _rows_of(database, secret_column) + ] + ) + + +_CIPHERTEXT: Final = TypeAdapter(str) + + +async def _secret_columns_in(database: SupportsRawQueries) -> tuple[_SecretColumn, ...]: + existing: Final = frozenset( + (row["table_name"], row["column_name"]) + for row in await database.query_raw( + "SELECT table_name, column_name FROM information_schema.columns " + "WHERE table_schema = ANY (current_schemas(false))" + ) + ) + return tuple( + secret_column for secret_column in _SECRET_COLUMNS if (secret_column.table, secret_column.column) in existing + ) + + +async def _rows_of(database: SupportsRawQueries, secret_column: _SecretColumn) -> tuple[Mapping[str, object], ...]: + marked_only: Final = ( + f" AND \"{secret_column.column}\"::text LIKE '%{CALLBACK_VAR_ENCRYPTED_PREFIX}%'" + if secret_column.only_rows_with_marked_ciphertexts + else "" + ) + return tuple( + await database.query_raw( + f'SELECT "{secret_column.primary_key}", "{secret_column.column}" FROM "{secret_column.table}" ' + f'WHERE "{secret_column.column}" IS NOT NULL{marked_only}' + ) + ) + + +async def _reencrypt_row( + database: SupportsRawQueries, + secret_column: _SecretColumn, + row: Mapping[str, object], + reencrypted: ReplaceCiphertext, +) -> int: + stored: Final = _STORED_VALUE.validate_python(row[secret_column.column]) + migrated, count = replace_ciphertexts(stored, reencrypted) + if count == 0: + return 0 + cast_to: Final = "::jsonb" if secret_column.is_json else "" + rows_updated: Final = await database.execute_raw( + f'UPDATE "{secret_column.table}" SET "{secret_column.column}" = $1{cast_to} ' + f'WHERE "{secret_column.primary_key}" = $2 AND "{secret_column.column}" = $3{cast_to}', + _as_sql_parameter(migrated, secret_column), + _PRIMARY_KEY.validate_python(row[secret_column.primary_key]), + _as_sql_parameter(stored, secret_column), + ) + return count if rows_updated else 0 + + +def _as_sql_parameter(value: JsonValue, secret_column: _SecretColumn) -> str: + return json.dumps(value) if secret_column.is_json else _CIPHERTEXT.validate_python(value) + + +class NothingToMigrate(Enum): + SALT_KEY_ENCRYPTS_STORED_VALUES = "salt_key_encrypts_stored_values" + NO_DATABASE = "no_database" + NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY = "nothing_encrypted_with_previous_key" + + +@dataclass(frozen=True, slots=True) +class Migrated: + migrated: int + remaining: int + + +@dataclass(frozen=True, slots=True) +class MigrationFailed: + error: Exception + + +MigrationOutcome = NothingToMigrate | Migrated | MigrationFailed + + +async def migrate_if_requested( + *, + environ: Mapping[str, str], + master_key: str | None, + connected_database: Callable[[], SupportsRawQueries | None], + log: Callable[[str], None], + raise_unless_tolerated: Callable[[Exception], None], +) -> MigrationOutcome | None: + previous_master_key: Final = environ.get(MIGRATE_FROM_MASTER_KEY_ENV_VAR) + if previous_master_key is None or master_key is None: + return None + outcome: Final = await migrate_from_previous_master_key( + previous_master_key=previous_master_key, + master_key=master_key, + salt_key_is_set=SALT_KEY_ENV_VAR in environ, + database=connected_database(), + log=log, + ) + if isinstance(outcome, MigrationFailed): + raise_unless_tolerated(outcome.error) + return outcome + + +async def migrate_from_previous_master_key( + *, + previous_master_key: str, + master_key: str, + salt_key_is_set: bool, + database: SupportsRawQueries | None, + log: Callable[[str], None], +) -> MigrationOutcome: + outcome: Final = await _migrate_or_failure( + previous_master_key=previous_master_key, + master_key=master_key, + salt_key_is_set=salt_key_is_set, + database=database, + log=log, + ) + log(describe_outcome(outcome)) + return outcome + + +async def _migrate_or_failure( + *, + previous_master_key: str, + master_key: str, + salt_key_is_set: bool, + database: SupportsRawQueries | None, + log: Callable[[str], None], +) -> MigrationOutcome: + try: + return await _migrate( + previous_master_key=previous_master_key, + master_key=master_key, + salt_key_is_set=salt_key_is_set, + database=database, + log=log, + ) + except Exception as error: # noqa: BLE001 # a value, so the boot applies its own database outage rule to it + return MigrationFailed(error=error) + + +async def _migrate( + *, + previous_master_key: str, + master_key: str, + salt_key_is_set: bool, + database: SupportsRawQueries | None, + log: Callable[[str], None], +) -> MigrationOutcome: + if salt_key_is_set: + return NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES + if database is None: + return NothingToMigrate.NO_DATABASE + if previous_master_key == master_key: + return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY + found: Final = await count_values_encrypted_with(database, previous_master_key) + if found == 0: + return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY + log(f"Re-encrypting {found} stored value(s) from the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} key to the new master key.") + migrated: Final = Migrated( + migrated=await reencrypt_stored_values(database, from_key=previous_master_key, to_key=master_key), + remaining=await count_values_encrypted_with(database, previous_master_key), + ) + another_worker_migrated_everything: Final = migrated == Migrated(migrated=0, remaining=0) + return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY if another_worker_migrated_everything else migrated + + +def describe_outcome(outcome: MigrationOutcome) -> str: + match outcome: + case NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES: + return ( + f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is set, but {SALT_KEY_ENV_VAR} is what encrypts your stored " + f"values, so there is nothing to migrate. You may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}." + ) + case NothingToMigrate.NO_DATABASE: + return ( + f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is set, but no database is connected, so nothing was migrated. If " + f"this proxy has no database, you may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}." + ) + case NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY: + return ( + f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is still set, but nothing in the database is left to migrate " + f"from that key. You may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}." + ) + case Migrated(migrated=migrated, remaining=0): + return ( + f"Done re-encrypting {migrated} stored value(s) with the new master key. You may now delete the " + f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} environment variable." + ) + case Migrated(migrated=migrated, remaining=remaining): + return ( + f"Re-encrypted {migrated} stored value(s), but {remaining} are still encrypted with the previous key " + f"because they changed during the migration. Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set and restart " + "the proxy to migrate them." + ) + case MigrationFailed(error=error): + cause: Final = f"{type(error).__name__}: {error}"[:300] + return ( + f"Could not migrate stored values from the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} key ({cause}). Values " + "still encrypted with the previous key cannot be read until the migration succeeds. Keep " + f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} set and restart the proxy once the database is reachable." + ) + case _: + assert_never(outcome) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index af25d418a63..e49d563328f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -342,6 +342,15 @@ from litellm.proxy.auth.login_throttle import ( warn_login_counters_are_per_worker, warn_source_login_limit_is_off, ) +from litellm.proxy.auth.master_key_boot_check import ( + MASTER_KEY_ENV_VAR, + SALT_KEY_ENV_VAR, + WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR, + announce_on_stderr_at_exit, + enforce_master_key_boot_verdict, + master_key_boot_verdict, + with_stored_secrets_counted, +) from litellm.proxy.auth.model_checks import ( expand_wildcard_deployments_for_model_info, get_all_fallbacks, @@ -464,6 +473,7 @@ from litellm.proxy.config_resolvers.settings_rules import ( ) from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router +from litellm.proxy.db.create_views import SupportsRawQueries from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( @@ -478,6 +488,10 @@ from litellm.proxy.db.gateway_request_tracking import ( GatewayRequestRedisBuffer, flush_gateway_requests, ) +from litellm.proxy.db.master_key_migration import ( + count_values_encrypted_with_or_none, + migrate_if_requested, +) from litellm.proxy.db.proxy_worker_heartbeat import ( PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS, ProxyWorkerHeartbeat, @@ -1123,6 +1137,14 @@ async def _initialize_shared_aiohttp_session(): return None +async def _connect_to_count_stored_values() -> SupportsRawQueries: + client: Final = prisma_client or PrismaClient( + database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj + ) + await client.connect() + return client.writer_db + + @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: global \ @@ -1215,6 +1237,22 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: if isinstance(worker_config, dict): await initialize(**worker_config) + enforce_master_key_boot_verdict( + await with_stored_secrets_counted( + master_key_boot_verdict( + master_key=master_key, + environment_master_key=os.getenv(MASTER_KEY_ENV_VAR), + general_settings=general_settings, + config_file_path=user_config_file_path, + override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True, + salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, + database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None, + ), + count_values_encrypted_with=partial(count_values_encrypted_with_or_none, _connect_to_count_stored_values), + ), + announce=announce_on_stderr_at_exit, + ) + # check if DATABASE_URL in environment - load from there if prisma_client is None: _db_url: Final[str | None] = get_secret("DATABASE_URL", None) @@ -1224,6 +1262,14 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: user_api_key_cache=user_api_key_cache, ) + await migrate_if_requested( + environ=os.environ, + master_key=master_key, + connected_database=lambda: None if prisma_client is None else prisma_client.writer_db, + log=verbose_proxy_logger.warning, + raise_unless_tolerated=PrismaDBExceptionHandler.handle_db_exception, + ) + if prisma_client is not None: async def _run_pw_migration(): diff --git a/render.yaml b/render.yaml index 18ad8ff2078..b61385c4d90 100644 --- a/render.yaml +++ b/render.yaml @@ -7,6 +7,8 @@ services: envVars: - key: PORT value: 4000 + - key: LITELLM_MASTER_KEY + generateValue: true numInstances: 1 healthCheckPath: /health/liveliness autoDeploy: true diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 3c6a6a58820..071191183df 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -55,6 +55,7 @@ IGNORE_FUNCTIONS = [ "apply_json_merge_patch", # max depth set (_MAX_MERGE_DEPTH=64); fails closed by raising ValueError at the cap. "_filter_argument_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by blocking the tool call at the cap. "_redact_scanned_content", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by returning "[REDACTED]" at the cap. + "replace_ciphertexts", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); walks stored JSON, which has no cycles, and leaves values below the cap untouched. "_iter_fallback_targets", # max depth set (2 * ROUTER_MAX_FALLBACKS); fails closed by raising ValueError at the cap. "_mergeable_branch", # max depth set (_MAX_SCHEMA_FLATTEN_DEPTH=32) plus a seen_refs cycle guard; passes the schema through untouched at the cap. "json_string_leaves", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); fails closed by raising at the cap so nothing goes unscanned. diff --git a/tests/e2e/ui/run_e2e.sh b/tests/e2e/ui/run_e2e.sh index beb1bc8bf3b..5c367ba0024 100755 --- a/tests/e2e/ui/run_e2e.sh +++ b/tests/e2e/ui/run_e2e.sh @@ -145,6 +145,7 @@ fi # --- Credentials --- export LITELLM_MASTER_KEY="sk-1234" +export LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY="true" export MOCK_LLM_URL="http://127.0.0.1:${MOCK_LLM_PORT}/v1" export E2E_MOCK_PRESIDIO_URL="http://127.0.0.1:${MOCK_PRESIDIO_PORT}" export DISABLE_SCHEMA_UPDATE="true" diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index 99c03b3438d..a730f6c10ee 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -54,6 +54,7 @@ def _clear_proxy_database_env() -> typing.Iterator[None]: # the LITELLM_MASTER_KEY env var, overriding whatever initialize() set from # the config file. We must set it here so the lifespan doesn't reset it to None. mp.setenv("LITELLM_MASTER_KEY", "sk-1234") + mp.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") try: yield finally: diff --git a/tests/proxy_behavior/management/conftest.py b/tests/proxy_behavior/management/conftest.py index 4c5b2ee9949..255b937bdd3 100644 --- a/tests/proxy_behavior/management/conftest.py +++ b/tests/proxy_behavior/management/conftest.py @@ -13,7 +13,7 @@ from prisma import Json from litellm.proxy.utils import hash_token -MASTER_KEY = "sk-1234" +MASTER_KEY = "sk-proxy-behavior-master-key" SCRATCH_PREFIX = "scratch-" diff --git a/tests/proxy_security_tests/test_master_key_not_in_db.py b/tests/proxy_security_tests/test_master_key_not_in_db.py index cb6e08d6746..2f00e57037a 100644 --- a/tests/proxy_security_tests/test_master_key_not_in_db.py +++ b/tests/proxy_security_tests/test_master_key_not_in_db.py @@ -20,8 +20,10 @@ def override_env_settings(monkeypatch): @pytest.fixture(scope="module") def test_client(): """Starting the test client triggers FastAPI startup, where Prisma connects to the DB.""" - with TestClient(app) as client: - yield client + with pytest.MonkeyPatch.context() as boot_env: + boot_env.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") + with TestClient(app) as client: + yield client @pytest.mark.asyncio diff --git a/tests/proxy_unit_tests/test_aproxy_startup.py b/tests/proxy_unit_tests/test_aproxy_startup.py index 98bf6ef8eb7..7604c96e07c 100644 --- a/tests/proxy_unit_tests/test_aproxy_startup.py +++ b/tests/proxy_unit_tests/test_aproxy_startup.py @@ -22,7 +22,7 @@ from litellm.proxy.proxy_server import ( @pytest.mark.asyncio -async def test_proxy_gunicorn_startup_direct_config(): +async def test_proxy_gunicorn_startup_direct_config(monkeypatch): """ gunicorn startup requires the config to be passed in via environment variables @@ -30,6 +30,7 @@ async def test_proxy_gunicorn_startup_direct_config(): Test both approaches """ + monkeypatch.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") try: from litellm._logging import verbose_proxy_logger, verbose_router_logger import logging @@ -59,7 +60,8 @@ async def test_proxy_gunicorn_startup_direct_config(): @pytest.mark.asyncio -async def test_proxy_gunicorn_startup_config_dict(): +async def test_proxy_gunicorn_startup_config_dict(monkeypatch): + monkeypatch.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") try: from litellm._logging import verbose_proxy_logger, verbose_router_logger import logging diff --git a/tests/test_litellm/proxy/auth/test_master_key_boot_check.py b/tests/test_litellm/proxy/auth/test_master_key_boot_check.py new file mode 100644 index 00000000000..aca65f4c7e6 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_master_key_boot_check.py @@ -0,0 +1,427 @@ +import asyncio +import re +import shutil +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path + +import pytest + +from litellm.proxy.auth.master_key_boot_check import ( + GENERATE_MASTER_KEY_COMMAND, + MASTER_KEY_ENV_VAR, + MIGRATE_FROM_MASTER_KEY_ENV_VAR, + PRINT_NEW_MASTER_KEY_COMMAND, + ROTATION_DOCS_URL, + WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR, + WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING, + ConfigFileSource, + EnvironmentSource, + MasterKeyBootVerdict, + SafeMasterKey, + StoredSecretsMigration, + UnsafeMasterKeyAllowed, + UnsafeMasterKeyError, + UnsafeMasterKeyReason, + UnsafeMasterKeyRefused, + announce_on_stderr_at_exit, + enforce_master_key_boot_verdict, + master_key_boot_verdict, + render_refusal, + with_stored_secrets_counted, +) + + +def _verdict( + master_key: str | None, + general_settings: Mapping[str, object] | None = None, + *, + environment_master_key: str | None = None, + config_file_path: str | None = None, + override_env_is_on: bool = False, + salt_key_is_set: bool = False, + database_is_configured: bool = False, +) -> MasterKeyBootVerdict: + return master_key_boot_verdict( + master_key=master_key, + environment_master_key=environment_master_key, + general_settings=general_settings or {}, + config_file_path=config_file_path, + override_env_is_on=override_env_is_on, + salt_key_is_set=salt_key_is_set, + database_is_configured=database_is_configured, + ) + + +@pytest.mark.parametrize( + ("master_key", "reason"), + [ + (None, UnsafeMasterKeyReason.NOT_SET), + ("", UnsafeMasterKeyReason.EMPTY), + (" \t\n", UnsafeMasterKeyReason.EMPTY), + ("sk-1234", UnsafeMasterKeyReason.PUBLICLY_KNOWN), + (" sk-1234\n", UnsafeMasterKeyReason.PUBLICLY_KNOWN), + ], +) +def test_unsafe_master_keys_are_refused_with_their_reason(master_key: str | None, reason: UnsafeMasterKeyReason): + verdict = _verdict(master_key) + + assert isinstance(verdict, UnsafeMasterKeyRefused) + assert verdict.reason is reason + + +@pytest.mark.parametrize("master_key", ["sk-12345", "sk-1234567890", "1234", "sk-qa-9f2c1e7a44b0d3"]) +def test_keys_that_only_resemble_the_known_default_are_safe(master_key: str): + assert _verdict(master_key) == SafeMasterKey() + + +@pytest.mark.parametrize("master_key", [None, "", "sk-1234"]) +def test_either_override_lets_an_unsafe_key_through(master_key: str | None): + from_env = _verdict(master_key, override_env_is_on=True) + from_yaml = _verdict(master_key, {WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING: True}) + + assert isinstance(from_env, UnsafeMasterKeyAllowed) + assert from_env == from_yaml + + +def test_override_switched_off_in_yaml_still_refuses(): + assert isinstance(_verdict("sk-1234", {WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING: False}), UnsafeMasterKeyRefused) + + +def test_yaml_master_key_is_the_source_even_when_it_resolved_to_nothing(): + verdict = _verdict(None, {"master_key": None}, config_file_path="/app/config.yaml") + + assert isinstance(verdict, UnsafeMasterKeyRefused) + assert verdict.source == ConfigFileSource(config_file_path="/app/config.yaml") + + +def test_yaml_master_key_is_the_source_when_it_differs_from_the_environment(): + verdict = _verdict( + "sk-1234", + {"master_key": "sk-1234"}, + environment_master_key="sk-qa-9f2c1e7a44b0d3", + config_file_path="/app/config.yaml", + ) + + assert isinstance(verdict, UnsafeMasterKeyRefused) + assert verdict.source == ConfigFileSource(config_file_path="/app/config.yaml") + + +@pytest.mark.parametrize("unsafe_key", ["sk-1234", ""]) +def test_environment_is_the_source_when_yaml_only_relays_the_environment_variable(unsafe_key: str): + verdict = _verdict( + unsafe_key, + {"master_key": unsafe_key}, + environment_master_key=unsafe_key, + config_file_path="/app/config.yaml", + ) + + assert isinstance(verdict, UnsafeMasterKeyRefused) + assert verdict.source == EnvironmentSource() + + +def test_environment_is_the_source_when_yaml_does_not_set_a_master_key(): + verdict = _verdict("sk-1234", {"database_url": "postgresql://db"}, config_file_path="/app/config.yaml") + + assert isinstance(verdict, UnsafeMasterKeyRefused) + assert verdict.source == EnvironmentSource() + + +@pytest.mark.parametrize( + ("master_key", "salt_key_is_set", "database_is_configured", "migration"), + [ + ("sk-1234", False, True, StoredSecretsMigration(from_master_key="sk-1234", encrypted_value_count=None)), + ("", False, True, StoredSecretsMigration(from_master_key="", encrypted_value_count=None)), + (" sk-1234\n", False, True, StoredSecretsMigration(from_master_key=" sk-1234\n", encrypted_value_count=None)), + ("sk-1234", True, True, None), + ("sk-1234", False, False, None), + (None, False, True, None), + ], +) +def test_migration_is_offered_from_the_exact_key_that_may_encrypt_a_database( + master_key: str | None, + salt_key_is_set: bool, + database_is_configured: bool, + migration: StoredSecretsMigration | None, +): + verdict = _verdict(master_key, salt_key_is_set=salt_key_is_set, database_is_configured=database_is_configured) + + assert isinstance(verdict, UnsafeMasterKeyRefused) + assert verdict.migration == migration + + +def _counted(verdict: MasterKeyBootVerdict, count: int | None) -> tuple[MasterKeyBootVerdict, list[str]]: + asked_about: list[str] = [] + + async def count_values_encrypted_with(signing_key: str) -> int | None: + asked_about.append(signing_key) + return count + + return asyncio.run(with_stored_secrets_counted(verdict, count_values_encrypted_with)), asked_about + + +def test_database_with_nothing_encrypted_needs_no_migration(): + counted, asked_about = _counted(_verdict("sk-1234", database_is_configured=True), 0) + + assert isinstance(counted, UnsafeMasterKeyRefused) + assert counted.migration is None + assert asked_about == ["sk-1234"] + + +@pytest.mark.parametrize("count", [4, None]) +def test_database_with_encrypted_values_or_unreadable_keeps_the_migration(count: int | None): + counted, _ = _counted(_verdict("", database_is_configured=True), count) + + assert isinstance(counted, UnsafeMasterKeyRefused) + assert counted.migration == StoredSecretsMigration(from_master_key="", encrypted_value_count=count) + + +@pytest.mark.parametrize( + "verdict", + [ + SafeMasterKey(), + UnsafeMasterKeyAllowed(reason=UnsafeMasterKeyReason.PUBLICLY_KNOWN), + _verdict("sk-1234", database_is_configured=False), + ], +) +def test_database_is_not_read_when_no_migration_is_on_the_table(verdict: MasterKeyBootVerdict): + counted, asked_about = _counted(verdict, 7) + + assert counted == verdict + assert asked_about == [] + + +def _refusal( + reason: UnsafeMasterKeyReason = UnsafeMasterKeyReason.PUBLICLY_KNOWN, + source: ConfigFileSource | EnvironmentSource = EnvironmentSource(), + environment_variable_is_set: bool = False, + migration: StoredSecretsMigration | None = None, +) -> UnsafeMasterKeyRefused: + return UnsafeMasterKeyRefused( + reason=reason, + source=source, + environment_variable_is_set=environment_variable_is_set, + migration=migration, + ) + + +_MIGRATION = StoredSecretsMigration(from_master_key="sk-1234", encrypted_value_count=3) + + +def test_config_refusal_names_the_file_and_tells_it_to_read_the_environment(): + text = render_refusal(_refusal(source=ConfigFileSource(config_file_path="/app/config.yaml"))) + + assert "general_settings.master_key in /app/config.yaml" in text + assert f"master_key: os.environ/{MASTER_KEY_ENV_VAR}" in text + assert GENERATE_MASTER_KEY_COMMAND in text + + +def test_environment_refusal_gives_the_command_without_a_config_step(): + text = render_refusal(_refusal(reason=UnsafeMasterKeyReason.NOT_SET, source=EnvironmentSource())) + + assert f"the {MASTER_KEY_ENV_VAR} environment variable" in text + assert GENERATE_MASTER_KEY_COMMAND in text + assert "os.environ/" not in text + + +@pytest.mark.parametrize( + ("master_key", "general_settings", "environment_master_key", "is_set"), + [ + (None, {}, None, False), + ("sk-1234", {"master_key": "sk-1234"}, None, False), + ("sk-1234", {}, "sk-1234", True), + ("sk-1234", {"master_key": "sk-1234"}, "", True), + ], +) +def test_refusal_records_whether_the_environment_variable_is_already_set( + master_key: str | None, general_settings: Mapping[str, object], environment_master_key: str | None, is_set: bool +): + refusal = _verdict(master_key, general_settings, environment_master_key=environment_master_key) + + assert isinstance(refusal, UnsafeMasterKeyRefused) + assert refusal.environment_variable_is_set is is_set + + +@pytest.mark.parametrize("source", [EnvironmentSource(), ConfigFileSource(config_file_path="/app/config.yaml")]) +def test_refusal_never_tells_a_user_with_an_exported_key_to_append_to_the_env_file( + source: ConfigFileSource | EnvironmentSource, +): + text = render_refusal(_refusal(source=source, environment_variable_is_set=True)) + + assert PRINT_NEW_MASTER_KEY_COMMAND in text + assert "tee" not in text + assert "wins over .env" in text + + +@pytest.mark.parametrize( + ("reason", "source", "source_line"), + [ + ( + UnsafeMasterKeyReason.NOT_SET, + EnvironmentSource(), + f"Neither general_settings.master_key nor the {MASTER_KEY_ENV_VAR} environment variable is set.", + ), + ( + UnsafeMasterKeyReason.PUBLICLY_KNOWN, + EnvironmentSource(), + f"It comes from the {MASTER_KEY_ENV_VAR} environment variable.", + ), + ( + UnsafeMasterKeyReason.NOT_SET, + ConfigFileSource(config_file_path="/app/config.yaml"), + "general_settings.master_key in /app/config.yaml is blank, " + "or points at an environment variable that is not set.", + ), + ( + UnsafeMasterKeyReason.EMPTY, + ConfigFileSource(config_file_path="/app/config.yaml"), + "It comes from general_settings.master_key in /app/config.yaml.", + ), + ], +) +def test_refusal_says_where_the_unsafe_key_came_from( + reason: UnsafeMasterKeyReason, source: ConfigFileSource | EnvironmentSource, source_line: str +): + assert render_refusal(_refusal(reason=reason, source=source)).splitlines()[1] == source_line + + +def test_migration_steps_appear_only_when_the_database_needs_them(): + with_migration = render_refusal(_refusal(migration=_MIGRATION)) + without_migration = render_refusal(_refusal(migration=None)) + + assert f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}=sk-1234" in with_migration + assert "holds 3 value(s) encrypted with this master key" in with_migration + assert ROTATION_DOCS_URL in with_migration + assert MIGRATE_FROM_MASTER_KEY_ENV_VAR not in without_migration + assert ROTATION_DOCS_URL not in without_migration + + +def test_unreadable_database_is_reported_as_unchecked_rather_than_counted(): + text = render_refusal( + _refusal(migration=StoredSecretsMigration(from_master_key="sk-1234", encrypted_value_count=None)) + ) + + assert "could not be checked" in text + assert "value(s)" not in text + assert f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}=sk-1234" in text + + +@pytest.mark.parametrize( + ("from_master_key", "assignment"), + [ + ("sk-1234", f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}=sk-1234"), + ("", f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}="), + (" sk-1234", f'{MIGRATE_FROM_MASTER_KEY_ENV_VAR}=" sk-1234"'), + ], +) +def test_migrate_from_assignment_carries_the_exact_previous_key(from_master_key: str, assignment: str): + text = render_refusal( + _refusal( + environment_variable_is_set=True, + migration=StoredSecretsMigration(from_master_key=from_master_key, encrypted_value_count=1), + ) + ) + + assert f" {assignment}\n" in text + + +def test_migration_with_an_exported_key_replaces_it_in_place_and_numbers_every_step(): + text = render_refusal( + _refusal( + source=ConfigFileSource(config_file_path="/app/config.yaml"), + environment_variable_is_set=True, + migration=_MIGRATION, + ) + ) + + assert "tee" not in text + assert PRINT_NEW_MASTER_KEY_COMMAND in text + assert [line[:2] for line in text.splitlines() if re.match(r"\d\. ", line)] == ["1.", "2.", "3.", "4."] + assert text.index("os.environ/") < text.index(MIGRATE_FROM_MASTER_KEY_ENV_VAR + "=") < text.index("Start the proxy") + + +@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") +@pytest.mark.parametrize("from_master_key", ["sk-1234", "", " sk-1234"]) +def test_printed_migration_commands_save_both_keys_to_the_env_file(tmp_path: Path, from_master_key: str): + from dotenv import dotenv_values + + text = render_refusal( + _refusal(migration=StoredSecretsMigration(from_master_key=from_master_key, encrypted_value_count=1)) + ) + commands = [line.strip() for line in text.splitlines() if line.strip().startswith("echo ")] + + subprocess.run(["bash", "-c", "\n".join(commands)], cwd=tmp_path, capture_output=True, text=True, check=True) + + saved = dotenv_values(tmp_path / ".env") + assert saved[MIGRATE_FROM_MASTER_KEY_ENV_VAR] == from_master_key + assert _verdict(saved[MASTER_KEY_ENV_VAR]) == SafeMasterKey() + assert len(commands) == 2 + + +@pytest.mark.parametrize("migration", [_MIGRATION, None]) +def test_override_hint_is_the_last_paragraph(migration: StoredSecretsMigration | None): + text = render_refusal(_refusal(migration=migration)) + last_paragraph = text.split("\n\n")[-1] + + assert WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR in last_paragraph + assert f"general_settings.{WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_SETTING}" in last_paragraph + + +@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") +def test_printed_command_saves_a_key_the_boot_check_accepts(tmp_path: Path): + completed = subprocess.run( + ["bash", "-c", GENERATE_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True + ) + + saved = (tmp_path / ".env").read_text() + match = re.fullmatch(rf"{MASTER_KEY_ENV_VAR}=(sk-[0-9a-f]{{64}})\n", saved) + assert match is not None + assert completed.stdout == saved + assert _verdict(match.group(1)) == SafeMasterKey() + + +@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl") +def test_rotation_command_prints_a_key_the_boot_check_accepts_and_saves_nothing(tmp_path: Path): + completed = subprocess.run( + ["bash", "-c", PRINT_NEW_MASTER_KEY_COMMAND], cwd=tmp_path, capture_output=True, text=True, check=True + ) + + assert re.fullmatch(r"sk-[0-9a-f]{64}\n", completed.stdout) is not None + assert _verdict(completed.stdout.strip()) == SafeMasterKey() + assert list(tmp_path.iterdir()) == [] + + +def test_refusal_announces_the_fix_and_aborts_the_boot(): + announced: list[str] = [] + refusal = _refusal(source=ConfigFileSource(config_file_path="/app/config.yaml")) + + with pytest.raises(UnsafeMasterKeyError, match="refused to start"): + enforce_master_key_boot_verdict(refusal, announce=announced.append) + + assert [message.strip() for message in announced] == [render_refusal(refusal)] + + +@pytest.mark.parametrize("verdict", [SafeMasterKey(), UnsafeMasterKeyAllowed(reason=UnsafeMasterKeyReason.NOT_SET)]) +def test_safe_and_overridden_keys_boot_without_announcing(verdict: MasterKeyBootVerdict): + announced: list[str] = [] + + enforce_master_key_boot_verdict(verdict, announce=announced.append) + + assert announced == [] + + +def test_announced_fix_is_the_last_thing_a_crashing_process_prints(): + crash_after_announcing = ( + "from litellm.proxy.auth.master_key_boot_check import announce_on_stderr_at_exit\n" + "announce_on_stderr_at_exit('THE FIX')\n" + "print('buffered stdout')\n" + "raise RuntimeError('lifespan failed')\n" + ) + + completed = subprocess.run([sys.executable, "-c", crash_after_announcing], capture_output=True, text=True) + + assert completed.returncode != 0 + assert "RuntimeError: lifespan failed" in completed.stderr + assert completed.stderr.endswith("THE FIX") + assert completed.stdout == "buffered stdout\n" diff --git a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py index 08cf1e45812..9c07242bd23 100644 --- a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -6,12 +6,16 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130 (nacl) ciphertext and new AES values coexist and decrypt correctly. """ +import base64 + import pytest from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import ( _V2_GCM_PREFIX, + decrypt_if_encrypted_with, decrypt_value_helper, + encrypt_value, encrypt_value_helper, ) @@ -185,3 +189,50 @@ def test_decrypt_failure_debug_log_omits_raw_value(monkeypatch): "the failing key should still be named in the breadcrumb" ) assert result == secret + + +@pytest.mark.parametrize("use_aes", [False, True]) +def test_explicit_key_decrypt_reads_only_values_written_under_that_key(monkeypatch, use_aes: bool): + if use_aes: + _use_aes(monkeypatch) + written_with_previous_key = encrypt_value_helper("stored-secret", new_encryption_key="sk-1234") + + assert decrypt_if_encrypted_with(written_with_previous_key, "sk-1234") == "stored-secret" + assert decrypt_if_encrypted_with(written_with_previous_key, "sk-another-key") is None + assert decrypt_value_helper(written_with_previous_key, key="t", exception_type="debug") is None + + +@pytest.mark.parametrize( + "not_a_ciphertext", + [ + "", + "gpt-5.4-mini", + "https://example.invalid/v1", + "v2:gcm:", + "aGVsbG8=", + "*", + "-", + "_", + "...", + " ", + "{}", + "[]", + "=", + ], +) +def test_explicit_key_decrypt_rejects_values_that_are_not_ciphertexts(not_a_ciphertext: str): + assert decrypt_if_encrypted_with(not_a_ciphertext, "sk-1234") is None + + +@pytest.mark.parametrize("use_aes", [False, True]) +def test_explicit_key_decrypt_tells_an_encrypted_empty_string_from_no_ciphertext(monkeypatch, use_aes: bool): + if use_aes: + _use_aes(monkeypatch) + + assert decrypt_if_encrypted_with(encrypt_value_helper("", new_encryption_key="sk-1234"), "sk-1234") == "" + + +def test_explicit_key_decrypt_supports_the_empty_master_key(): + written_with_empty_key = encrypt_value(value="stored-secret", signing_key="") + + assert decrypt_if_encrypted_with(base64.urlsafe_b64encode(written_with_empty_key).decode(), "") == "stored-secret" diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py index 65e12b7d777..8ef5017a952 100644 --- a/tests/test_litellm/proxy/conftest.py +++ b/tests/test_litellm/proxy/conftest.py @@ -268,6 +268,7 @@ def create_proxy_test_client( # Set environment variables set_proxy_environment_variables(monkeypatch, database_url=database_url) + monkeypatch.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") # Initialize proxy asyncio.run(initialize(config=config_fp, debug=init_options.get("debug", False))) diff --git a/tests/test_litellm/proxy/db/test_master_key_migration.py b/tests/test_litellm/proxy/db/test_master_key_migration.py new file mode 100644 index 00000000000..9c0fc163b9f --- /dev/null +++ b/tests/test_litellm/proxy/db/test_master_key_migration.py @@ -0,0 +1,563 @@ +import json +import re +from collections.abc import Mapping, Sequence +from functools import reduce + +import pytest +from pydantic import JsonValue + +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.proxy import proxy_server +from litellm.proxy.auth.master_key_boot_check import MIGRATE_FROM_MASTER_KEY_ENV_VAR, SALT_KEY_ENV_VAR +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_value_helper +from litellm.proxy.db.master_key_migration import ( + _SECRET_COLUMNS, + Migrated, + MigrationFailed, + NothingToMigrate, + count_values_encrypted_with, + describe_outcome, + migrate_from_previous_master_key, + migrate_if_requested, + reencrypt_stored_values, + replace_ciphertexts, +) + +PREVIOUS_KEY = "sk-1234" +NEW_KEY = "sk-qa-9f2c1e7a44b0d3" +UNRELATED_KEY = "sk-some-other-deployment" + +Tables = dict[str, list[dict[str, object]]] + + +@pytest.fixture(autouse=True) +def _legacy_algorithm_and_no_salt_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv(SALT_KEY_ENV_VAR, raising=False) + monkeypatch.setattr(proxy_server, "general_settings", {}) + + +def _encrypted(plaintext: str, key: str = PREVIOUS_KEY) -> str: + return str(encrypt_value_helper(plaintext, new_encryption_key=key)) + + +class _FakeDatabase: + def __init__(self, tables: Tables, tables_missing_from_the_schema: frozenset[str] = frozenset()) -> None: + self.tables = tables + self.tables_missing_from_the_schema = tables_missing_from_the_schema + self.writes: list[tuple[str, str, str]] = [] + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + if "information_schema.columns" in query: + assert "table_schema = ANY (current_schemas(false))" in query + return [ + {"table_name": secret_column.table, "column_name": secret_column.column} + for secret_column in _SECRET_COLUMNS + if secret_column.table not in self.tables_missing_from_the_schema + ] + select = re.fullmatch(r'SELECT "(\w+)", "(\w+)" FROM "(\w+)" WHERE "\2" IS NOT NULL(.*)', query) + assert select is not None, query + primary_key, column, table, row_filter = select.groups() + assert row_filter in ("", f" AND \"{column}\"::text LIKE '%litellm_enc::%'") + assert table not in self.tables_missing_from_the_schema, f'relation "{table}" does not exist' + return [ + {primary_key: row[primary_key], column: json.loads(json.dumps(row[column]))} + for row in self.tables.get(table, []) + if row.get(column) is not None and (not row_filter or "litellm_enc::" in json.dumps(row[column])) + ] + + async def execute_raw(self, query: str, *args: object) -> int: + update = re.fullmatch(r'UPDATE "(\w+)" SET "(\w+)" = \$1(::jsonb|) WHERE "(\w+)" = \$2 AND "\2" = \$3\3', query) + assert update is not None, query + table, column, json_cast, primary_key = update.groups() + new_value, row_id, expected = ( + json.loads(str(arg)) if json_cast and index != 1 else arg for index, arg in enumerate(args) + ) + matching = [row for row in self.tables[table] if row[primary_key] == row_id and row[column] == expected] + for row in matching: + row[column] = new_value + self.writes.append((table, column, str(row_id))) + return len(matching) + + +class _DatabaseThatMustNotBeTouched: + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + raise AssertionError(f"unexpected read: {query}") + + async def execute_raw(self, query: str, *args: object) -> int: + raise AssertionError(f"unexpected write: {query}") + + +def _seeded_tables() -> Tables: + return { + "LiteLLM_ProxyModelTable": [ + { + "model_id": "model-1", + "litellm_params": { + "api_key": _encrypted("provider-key"), + "model": _encrypted("openai/gpt-5.4-mini"), + "rpm": 10, + "use_in_pass_through": False, + "api_base": None, + }, + }, + {"model_id": "model-2", "litellm_params": {"api_key": _encrypted("other-deployment", UNRELATED_KEY)}}, + ], + "LiteLLM_Config": [ + {"param_name": "environment_variables", "param_value": {"LANGFUSE_SECRET_KEY": _encrypted("env-secret")}}, + {"param_name": "general_settings", "param_value": {"proxy_batch_write_at": 10, "ui_name": "plain text"}}, + {"param_name": "cleared", "param_value": None}, + ], + "LiteLLM_MCPServerTable": [ + { + "server_id": "mcp-1", + "credentials": {"auth_value": _encrypted("mcp-token"), "aws_region_name": "us-east-1"}, + "static_headers": _encrypted('{"X-Api-Key": "header-secret"}'), + "env_vars": [ + {"name": "GLOBAL", "scope": "global", "value": _encrypted("global-env")}, + {"name": "PER_USER", "scope": "user", "value": ""}, + ], + "env": {}, + } + ], + "LiteLLM_MCPUserCredentials": [{"id": "cred-row-1", "credential_b64": _encrypted("byok-secret")}], + "LiteLLM_TeamTable": [ + { + "team_id": "team-1", + "metadata": { + "logging": [ + { + "callback_name": "langfuse", + "callback_vars": { + "langfuse_host": "https://example.invalid", + "langfuse_secret_key": "litellm_enc::" + _encrypted("team-callback-secret"), + }, + } + ] + }, + }, + {"team_id": "team-without-callbacks", "metadata": {"note": _encrypted("unmarked, so never selected")}}, + ], + } + + +_VALUES_UNDER_THE_PREVIOUS_KEY = 8 + + +@pytest.mark.asyncio +async def test_reencryption_moves_every_stored_shape_to_the_new_key_and_nothing_else(): + tables = _seeded_tables() + untouched_before = json.dumps( + [tables["LiteLLM_ProxyModelTable"][1], tables["LiteLLM_Config"][1:], tables["LiteLLM_TeamTable"][1]] + ) + + migrated = await reencrypt_stored_values(_FakeDatabase(tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY) + + assert migrated == _VALUES_UNDER_THE_PREVIOUS_KEY + model_params = tables["LiteLLM_ProxyModelTable"][0]["litellm_params"] + assert isinstance(model_params, dict) + assert decrypt_if_encrypted_with(model_params["api_key"], NEW_KEY) == "provider-key" + 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) + 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" + assert decrypt_if_encrypted_with(mcp_server["env_vars"][0]["value"], NEW_KEY) == "global-env" + assert mcp_server["env_vars"][1] == {"name": "PER_USER", "scope": "user", "value": ""} + assert ( + decrypt_if_encrypted_with(tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"], NEW_KEY) == "byok-secret" + ) + callback_secret = tables["LiteLLM_TeamTable"][0]["metadata"]["logging"][0]["callback_vars"]["langfuse_secret_key"] + assert callback_secret.startswith("litellm_enc::") + assert decrypt_if_encrypted_with(callback_secret.removeprefix("litellm_enc::"), NEW_KEY) == "team-callback-secret" + assert untouched_before == json.dumps( + [tables["LiteLLM_ProxyModelTable"][1], tables["LiteLLM_Config"][1:], tables["LiteLLM_TeamTable"][1]] + ) + + +@pytest.mark.asyncio +async def test_count_follows_the_values_from_the_previous_key_to_the_new_one(): + database = _FakeDatabase(_seeded_tables()) + + assert await count_values_encrypted_with(database, PREVIOUS_KEY) == _VALUES_UNDER_THE_PREVIOUS_KEY + assert await count_values_encrypted_with(database, NEW_KEY) == 0 + + await reencrypt_stored_values(database, from_key=PREVIOUS_KEY, to_key=NEW_KEY) + + assert await count_values_encrypted_with(database, PREVIOUS_KEY) == 0 + assert await count_values_encrypted_with(database, NEW_KEY) == _VALUES_UNDER_THE_PREVIOUS_KEY + assert await count_values_encrypted_with(database, UNRELATED_KEY) == 1 + + +@pytest.mark.asyncio +async def test_only_rows_holding_values_under_the_previous_key_are_written(): + database = _FakeDatabase(_seeded_tables()) + + await reencrypt_stored_values(database, from_key=PREVIOUS_KEY, to_key=NEW_KEY) + + assert sorted(database.writes) == [ + ("LiteLLM_Config", "param_value", "environment_variables"), + ("LiteLLM_MCPServerTable", "credentials", "mcp-1"), + ("LiteLLM_MCPServerTable", "env_vars", "mcp-1"), + ("LiteLLM_MCPServerTable", "static_headers", "mcp-1"), + ("LiteLLM_MCPUserCredentials", "credential_b64", "cred-row-1"), + ("LiteLLM_ProxyModelTable", "litellm_params", "model-1"), + ("LiteLLM_TeamTable", "metadata", "team-1"), + ] + + +@pytest.mark.asyncio +async def test_plaintext_that_base64_decodes_to_nothing_is_neither_counted_nor_rewritten(): + settings = {"allowed_routes": ["*"], "ui_name": "-", "separator": "...", "blank": " ", "shape": "{}"} + tables: Tables = { + "LiteLLM_Config": [{"param_name": "general_settings", "param_value": dict(settings)}], + "LiteLLM_VerificationToken": [ + {"token": "hashed", "metadata": {"notes": "...", "secret": "litellm_enc::" + _encrypted("callback-secret")}} + ], + } + database = _FakeDatabase(tables) + + found = await count_values_encrypted_with(database, PREVIOUS_KEY) + migrated = await reencrypt_stored_values(database, from_key=PREVIOUS_KEY, to_key=NEW_KEY) + + assert found == migrated == 1 + assert tables["LiteLLM_Config"][0]["param_value"] == settings + assert tables["LiteLLM_VerificationToken"][0]["metadata"]["notes"] == "..." + assert database.writes == [("LiteLLM_VerificationToken", "metadata", "hashed")] + + +@pytest.mark.asyncio +async def test_schema_without_some_of_the_tables_is_migrated_for_the_tables_it_has(): + missing = frozenset({"LiteLLM_MCPUserCredentials", "LiteLLM_SSOIdentityAssertion"}) + tables = _seeded_tables() + database = _FakeDatabase(tables, tables_missing_from_the_schema=missing) + + found = await count_values_encrypted_with(database, PREVIOUS_KEY) + migrated = await reencrypt_stored_values(database, from_key=PREVIOUS_KEY, to_key=NEW_KEY) + + assert found == migrated == _VALUES_UNDER_THE_PREVIOUS_KEY - 1 + assert decrypt_if_encrypted_with(str(tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"]), PREVIOUS_KEY) + + +class _SomeoneEditsEachRowAfterItIsRead(_FakeDatabase): + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + rows = await super().query_raw(query, *args) + for row in self.tables.get("LiteLLM_MCPUserCredentials", []): + row["credential_b64"] = "edited-by-an-admin" + return rows + + +@pytest.mark.asyncio +async def test_value_edited_while_the_migration_runs_is_not_overwritten(): + tables: Tables = {"LiteLLM_MCPUserCredentials": [{"id": "cred-row-1", "credential_b64": _encrypted("byok-secret")}]} + + migrated = await reencrypt_stored_values( + _SomeoneEditsEachRowAfterItIsRead(tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY + ) + + assert migrated == 0 + assert tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"] == "edited-by-an-admin" + + +def test_replacing_ciphertexts_keeps_structure_markers_and_non_strings(): + value = {"keep": [1, True, None, "plain"], "swap": ["old", {"nested": "litellm_enc::old"}]} + + replaced, count = replace_ciphertexts(value, lambda text: "new" if text == "old" else None) + + assert replaced == {"keep": [1, True, None, "plain"], "swap": ["new", {"nested": "litellm_enc::new"}]} + assert count == 2 + assert value["swap"] == ["old", {"nested": "litellm_enc::old"}] + + +def _nested(levels: int, leaf: str) -> JsonValue: + return reduce(lambda inner, _: [inner], range(levels), leaf) + + +@pytest.mark.parametrize("levels_past_the_cap, replaced_count", [(0, 1), (1, 0), (50, 0)]) +def test_walk_stops_at_the_recursion_cap_and_leaves_deeper_values_as_they_were( + levels_past_the_cap: int, replaced_count: int +): + value = _nested(DEFAULT_MAX_RECURSE_DEPTH + levels_past_the_cap, "old") + + replaced, count = replace_ciphertexts(value, lambda text: "new") + + assert count == replaced_count + assert replaced == _nested(DEFAULT_MAX_RECURSE_DEPTH + levels_past_the_cap, "new" if replaced_count else "old") + + +async def _run( + database: _FakeDatabase | _DatabaseThatMustNotBeTouched | None, + *, + previous_master_key: str = PREVIOUS_KEY, + master_key: str = NEW_KEY, + salt_key_is_set: bool = False, +) -> tuple[object, list[str]]: + logged: list[str] = [] + outcome = await migrate_from_previous_master_key( + previous_master_key=previous_master_key, + master_key=master_key, + salt_key_is_set=salt_key_is_set, + database=database, + log=logged.append, + ) + return outcome, logged + + +@pytest.mark.asyncio +async def test_migration_announces_itself_then_says_the_variable_can_go(): + database = _FakeDatabase(_seeded_tables()) + + outcome, logged = await _run(database) + + assert outcome == Migrated(migrated=_VALUES_UNDER_THE_PREVIOUS_KEY, remaining=0) + assert len(logged) == 2 + assert logged[0].startswith(f"Re-encrypting {_VALUES_UNDER_THE_PREVIOUS_KEY} stored value(s)") + assert logged[1].startswith(f"Done re-encrypting {_VALUES_UNDER_THE_PREVIOUS_KEY} stored value(s)") + assert f"You may now delete the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} environment variable" in logged[1] + + +@pytest.mark.asyncio +async def test_variable_left_set_after_the_migration_is_a_no_op_with_one_notice(): + database = _FakeDatabase(_seeded_tables()) + await _run(database) + writes_after_the_migration = list(database.writes) + stored_after_the_migration = json.dumps(database.tables) + + outcome, logged = await _run(database) + + assert outcome is NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY + assert logged == [describe_outcome(NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY)] + assert database.writes == writes_after_the_migration + assert json.dumps(database.tables) == stored_after_the_migration + + +@pytest.mark.asyncio +async def test_empty_previous_master_key_migrates_like_any_other(): + tables: Tables = {"LiteLLM_CredentialsTable": [{"credential_id": "c1", "credential_values": {"api_key": "x"}}]} + tables["LiteLLM_CredentialsTable"][0]["credential_values"] = {"api_key": _encrypted_with_empty_key("cred-secret")} + + outcome, _ = await _run(_FakeDatabase(tables), previous_master_key="") + + assert outcome == Migrated(migrated=1, remaining=0) + stored = tables["LiteLLM_CredentialsTable"][0]["credential_values"] + assert isinstance(stored, dict) + assert decrypt_if_encrypted_with(stored["api_key"], NEW_KEY) == "cred-secret" + + +def _encrypted_with_empty_key(plaintext: str) -> str: + import base64 + + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value + + return base64.urlsafe_b64encode(encrypt_value(value=plaintext, signing_key="")).decode() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("salt_key_is_set", "has_database", "master_key", "expected"), + [ + (True, True, NEW_KEY, NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES), + (False, False, NEW_KEY, NothingToMigrate.NO_DATABASE), + (False, True, PREVIOUS_KEY, NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY), + ], +) +async def test_database_is_left_alone_when_the_master_key_cannot_have_encrypted_it( + salt_key_is_set: bool, has_database: bool, master_key: str, expected: NothingToMigrate +): + outcome, logged = await _run( + _DatabaseThatMustNotBeTouched() if has_database else None, + master_key=master_key, + salt_key_is_set=salt_key_is_set, + ) + + assert outcome is expected + assert logged == [describe_outcome(expected)] + + +class _AnotherWorkerMigratesRightAfterTheFirstCount(_FakeDatabase): + def __init__(self, tables: Tables) -> None: + super().__init__(tables) + self.reads = 0 + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + rows = await super().query_raw(query, *args) + self.reads += 1 + if self.reads == len(_SECRET_COLUMNS): + await reencrypt_stored_values(_FakeDatabase(self.tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY) + return rows + + +@pytest.mark.asyncio +async def test_worker_that_loses_the_race_reports_nothing_left_instead_of_zero_values_done(): + database = _AnotherWorkerMigratesRightAfterTheFirstCount(_seeded_tables()) + + outcome, logged = await _run(database) + + assert outcome is NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY + assert database.writes == [] + assert logged[-1] == describe_outcome(NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY) + + +@pytest.mark.asyncio +async def test_values_that_could_not_be_written_keep_the_variable_in_place(): + tables: Tables = {"LiteLLM_MCPUserCredentials": [{"id": "cred-row-1", "credential_b64": _encrypted("byok-secret")}]} + + class _EveryWriteLosesToACurrentEdit(_FakeDatabase): + async def execute_raw(self, query: str, *args: object) -> int: + return 0 + + outcome, logged = await _run(_EveryWriteLosesToACurrentEdit(tables)) + + assert outcome == Migrated(migrated=0, remaining=1) + assert "1 are still encrypted with the previous key" in logged[-1] + assert f"Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set" in logged[-1] + assert "You may now delete" not in logged[-1] + + +@pytest.mark.parametrize( + "outcome", [*NothingToMigrate, Migrated(migrated=5, remaining=0)], ids=lambda outcome: str(outcome) +) +def test_every_finished_outcome_tells_the_user_the_variable_can_be_deleted(outcome: NothingToMigrate | Migrated): + message = describe_outcome(outcome) + + assert "ou may now delete" in message + assert MIGRATE_FROM_MASTER_KEY_ENV_VAR in message + + +def test_salt_key_outcome_names_the_salt_key_as_the_reason(): + assert SALT_KEY_ENV_VAR in describe_outcome(NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES) + assert SALT_KEY_ENV_VAR not in describe_outcome(NothingToMigrate.NO_DATABASE) + + +@pytest.mark.asyncio +async def test_encrypted_empty_string_is_migrated_like_any_other_value(): + tables: Tables = {"LiteLLM_MCPUserCredentials": [{"id": "cred-row-1", "credential_b64": _encrypted("")}]} + + migrated = await reencrypt_stored_values(_FakeDatabase(tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY) + + assert migrated == 1 + assert decrypt_if_encrypted_with(str(tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"]), NEW_KEY) == "" + + +class _DatabaseIsDown(_FakeDatabase): + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + raise ConnectionError("Can't reach database server") + + +@pytest.mark.asyncio +async def test_database_error_during_the_migration_comes_back_as_a_value_and_is_logged(): + outcome, logged = await _run(_DatabaseIsDown(_seeded_tables())) + + assert isinstance(outcome, MigrationFailed) + assert isinstance(outcome.error, ConnectionError) + assert len(logged) == 1 + assert "ConnectionError: Can't reach database server" in logged[0] + assert f"Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set" in logged[0] + assert "ou may now delete" not in logged[0] + + +def _raise(error: Exception) -> None: + raise error + + +def _tolerate(error: Exception) -> None: + return None + + +@pytest.mark.asyncio +async def test_boot_stops_on_a_failed_migration_when_the_outage_is_not_tolerated(): + logged: list[str] = [] + + with pytest.raises(ConnectionError, match="Can't reach database server"): + await migrate_if_requested( + environ={MIGRATE_FROM_MASTER_KEY_ENV_VAR: PREVIOUS_KEY}, + master_key=NEW_KEY, + connected_database=lambda: _DatabaseIsDown(_seeded_tables()), + log=logged.append, + raise_unless_tolerated=_raise, + ) + + assert len(logged) == 1 + assert f"Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set" in logged[0] + + +@pytest.mark.asyncio +async def test_boot_continues_past_a_failed_migration_when_the_outage_is_tolerated(): + outcome = await migrate_if_requested( + environ={MIGRATE_FROM_MASTER_KEY_ENV_VAR: PREVIOUS_KEY}, + master_key=NEW_KEY, + connected_database=lambda: _DatabaseIsDown(_seeded_tables()), + log=lambda line: None, + raise_unless_tolerated=_tolerate, + ) + + assert isinstance(outcome, MigrationFailed) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("previous_master_key", [PREVIOUS_KEY, ""]) +async def test_boot_migrates_from_the_environment_variable_to_the_running_master_key(previous_master_key: str): + tables: Tables = { + "LiteLLM_CredentialsTable": [ + { + "credential_id": "cred-1", + "credential_values": { + "api_key": _encrypted("sk-provider", previous_master_key) + if previous_master_key + else _encrypted_with_empty_key("sk-provider") + }, + } + ] + } + logged: list[str] = [] + + outcome = await migrate_if_requested( + environ={MIGRATE_FROM_MASTER_KEY_ENV_VAR: previous_master_key}, + master_key=NEW_KEY, + connected_database=lambda: _FakeDatabase(tables), + log=logged.append, + raise_unless_tolerated=_raise, + ) + + assert outcome == Migrated(migrated=1, remaining=0) + stored = tables["LiteLLM_CredentialsTable"][0]["credential_values"] + assert isinstance(stored, dict) + assert decrypt_if_encrypted_with(stored["api_key"], NEW_KEY) == "sk-provider" + assert "Done re-encrypting 1 stored value(s)" in logged[-1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "environ, master_key, outcome", + [ + ({}, NEW_KEY, None), + ({MIGRATE_FROM_MASTER_KEY_ENV_VAR: PREVIOUS_KEY}, None, None), + ( + {MIGRATE_FROM_MASTER_KEY_ENV_VAR: PREVIOUS_KEY, SALT_KEY_ENV_VAR: "a-salt-key"}, + NEW_KEY, + NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES, + ), + ], + ids=["variable-not-set", "no-master-key", "salt-key-set"], +) +async def test_boot_leaves_the_database_alone_unless_a_migration_was_requested_and_can_apply( + environ: dict[str, str], master_key: str | None, outcome: NothingToMigrate | None +): + logged: list[str] = [] + database_handles_taken: list[str] = [] + + def connected_database() -> _DatabaseThatMustNotBeTouched: + database_handles_taken.append("taken") + return _DatabaseThatMustNotBeTouched() + + result = await migrate_if_requested( + environ=environ, + master_key=master_key, + connected_database=connected_database, + log=logged.append, + raise_unless_tolerated=_raise, + ) + + assert result is outcome + assert len(database_handles_taken) == (0 if outcome is None else 1) + assert len(logged) == (0 if outcome is None else 1) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 08a4621de24..468180ffe52 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1628,6 +1628,188 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): assert master_key == test_resolved_key +def _boot_with_general_settings(monkeypatch, tmp_path, general_settings): + import yaml + + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"general_settings": general_settings})) + for name in ( + "LITELLM_MASTER_KEY", + "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", + "LITELLM_MIGRATE_FROM_MASTER_KEY", + "LITELLM_SALT_KEY", + "WORKER_CONFIG", + "DATABASE_URL", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path)) + scheduler_left_on_a_closed_event_loop_by_an_earlier_test = "litellm.proxy.proxy_server.scheduler" + monkeypatch.setattr(scheduler_left_on_a_closed_event_loop_by_an_earlier_test, None) + announced = [] + monkeypatch.setattr("litellm.proxy.proxy_server.announce_on_stderr_at_exit", announced.append) + return config_path, announced + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings", + [{"master_key": "sk-1234"}, {"master_key": ""}, {"master_key": None}, {}], + ids=["publicly-known", "empty", "yaml-null", "no-general-settings"], +) +async def test_proxy_startup_refuses_an_unsafe_master_key_even_when_the_database_is_unreachable( + monkeypatch, tmp_path, general_settings +): + from fastapi import FastAPI + + from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError + from litellm.proxy.proxy_server import proxy_startup_event + + async def unreachable(): + raise ConnectionError("database is down") + + _, announced = _boot_with_general_settings(monkeypatch, tmp_path, general_settings) + monkeypatch.setenv("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable") + monkeypatch.setattr("litellm.proxy.proxy_server._connect_to_count_stored_values", unreachable) + + with pytest.raises(UnsafeMasterKeyError): + async with proxy_startup_event(FastAPI()): + pass + + assert len(announced) == 1 + assert "sk-$(openssl rand -hex 32)" in announced[0] + key_can_have_encrypted_the_database = general_settings.get("master_key") is not None + assert ("could not be checked" in announced[0]) == key_can_have_encrypted_the_database + + +class _DatabaseWithOneStoredCredential: + def __init__(self, ciphertext): + self._ciphertext = ciphertext + + async def query_raw(self, query, *args): + if "information_schema.columns" in query: + return [{"table_name": "LiteLLM_CredentialsTable", "column_name": "credential_values"}] + return [{"credential_id": "cred-1", "credential_values": {"api_key": self._ciphertext}}] + + async def execute_raw(self, query, *args): + raise AssertionError("a refused boot must not write to the database") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("encrypted_with, asks_to_migrate", [("sk-1234", True), ("sk-some-other-key", False)]) +async def test_proxy_startup_asks_to_migrate_only_when_the_database_holds_values_under_the_unsafe_key( + monkeypatch, tmp_path, encrypted_with, asks_to_migrate +): + from fastapi import FastAPI + + from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.proxy_server import proxy_startup_event + + database = _DatabaseWithOneStoredCredential(encrypt_value_helper("sk-provider", new_encryption_key=encrypted_with)) + + async def connected(): + return database + + _, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-1234"}) + monkeypatch.setenv("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable") + monkeypatch.setattr("litellm.proxy.proxy_server._connect_to_count_stored_values", connected) + + with pytest.raises(UnsafeMasterKeyError): + async with proxy_startup_event(FastAPI()): + pass + + assert ("LITELLM_MIGRATE_FROM_MASTER_KEY=sk-1234" in announced[0]) == asks_to_migrate + assert ("holds 1 value(s) encrypted with this master key" in announced[0]) == asks_to_migrate + + +@pytest.mark.asyncio +async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_deleted(monkeypatch, tmp_path, caplog): + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + async with proxy_startup_event(FastAPI()): + pass + + notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message] + assert len(notices) == 1 + assert "you may now delete LITELLM_MIGRATE_FROM_MASTER_KEY" in notices[0] + + +class _PrismaClientWhoseDatabaseRejectsQueries: + class _Database: + async def query_raw(self, query, *args): + raise RuntimeError("permission denied for table LiteLLM_CredentialsTable") + + writer_db = _Database() + + +@pytest.mark.asyncio +async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatch, tmp_path, caplog): + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "sk-1234") + + with ( + caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"), + pytest.raises(RuntimeError, match="permission denied"), + ): + async with proxy_startup_event(FastAPI()): + pass + + notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message] + assert len(notices) == 1 + assert "Could not migrate stored values" in notices[0] + + +@pytest.mark.asyncio +async def test_proxy_startup_names_the_config_file_that_set_the_unsafe_key(monkeypatch, tmp_path): + from fastapi import FastAPI + + from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError + from litellm.proxy.proxy_server import proxy_startup_event + + config_path, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-1234"}) + + with pytest.raises(UnsafeMasterKeyError): + async with proxy_startup_event(FastAPI()): + pass + + assert str(config_path) in announced[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("override", ["yaml", "env"]) +async def test_proxy_startup_boots_an_unsafe_master_key_under_the_override(monkeypatch, tmp_path, override): + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + general_settings = { + "master_key": "sk-1234", + **({"dangerously_permit_weak_or_unset_master_key": True} if override == "yaml" else {}), + } + _, announced = _boot_with_general_settings(monkeypatch, tmp_path, general_settings) + if override == "env": + monkeypatch.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") + + async with proxy_startup_event(FastAPI()): + from litellm.proxy.proxy_server import master_key + + assert master_key == "sk-1234" + + assert announced == [] + + def test_team_info_masking(): """ Test that sensitive team information is properly masked diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 57ff2fe9d87..1349dd0d099 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26830,6 +26830,11 @@ export interface components { * @description override user_api_key_auth with your own auth script - https://docs.litellm.ai/docs/proxy/virtual_keys#custom-auth */ custom_auth?: string | null; + /** + * Dangerously Permit Weak Or Unset Master Key + * @description local development only: start even when master_key is unset, empty, or a publicly known default + */ + dangerously_permit_weak_or_unset_master_key?: boolean | null; /** @description custom args for instantiating dynamodb client - e.g. billing provision */ database_args?: components["schemas"]["DynamoDBArgs"] | null; /**