diff --git a/litellm/proxy/auth/master_key_boot_check.py b/litellm/proxy/auth/master_key_boot_check.py index 68cbf2ea6e1..c8c2ea3539c 100644 --- a/litellm/proxy/auth/master_key_boot_check.py +++ b/litellm/proxy/auth/master_key_boot_check.py @@ -1,7 +1,7 @@ import atexit import sys -from collections.abc import Callable, Mapping -from dataclasses import dataclass +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, replace from enum import Enum from types import MappingProxyType from typing import Final @@ -15,6 +15,7 @@ UNSAFE_PROXY_OVERRIDE_ENV_VAR: Final = "LITELLM_DANGEROUSLY_ALLOW_UNSAFE_PROXY" 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)" @@ -51,12 +52,18 @@ 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 - stored_credentials_need_rotation: bool + migration: StoredSecretsMigration | None MasterKeyBootVerdict = SafeMasterKey | UnsafeMasterKeyAllowed | UnsafeMasterKeyRefused @@ -90,12 +97,23 @@ def master_key_boot_verdict( else EnvironmentSource() ), environment_variable_is_set=environment_master_key is not None, - stored_credentials_need_rotation=( - reason is UnsafeMasterKeyReason.PUBLICLY_KNOWN and not salt_key_is_set and database_is_configured + 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(): @@ -129,7 +147,7 @@ def render_refusal(refusal: UnsafeMasterKeyRefused) -> str: return "\n\n".join( ( f"LiteLLM proxy refused to start: {_REFUSAL_HEADLINE[refusal.reason]}\n{_source_line(refusal)}", - _ROTATE_INSTEAD_OF_REPLACING if refusal.stored_credentials_need_rotation else _fix_steps(refusal), + _fix_steps(refusal), _OVERRIDE_HINT, ) ) @@ -168,12 +186,9 @@ _REPLACE_EXPORTED_KEY_STEP: Final = ( " already exported in the environment wins over .env." ) -_ROTATE_INSTEAD_OF_REPLACING: Final = ( - f"Credentials stored in your database are encrypted with this master key because {SALT_KEY_ENV_VAR} is not\n" - "set, so replacing the key makes them undecryptable. Rotate it by following this guide, which re-encrypts them:\n" - f" {ROTATION_DOCS_URL}\n" - "Generate the new key for it with (save it only once the guide says to):\n" - f" {PRINT_NEW_MASTER_KEY_COMMAND}" +_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 = ( @@ -218,15 +233,59 @@ def _source_line(refusal: UnsafeMasterKeyRefused) -> str: def _fix_steps(refusal: UnsafeMasterKeyRefused) -> str: - set_key_step: Final = _REPLACE_EXPORTED_KEY_STEP if refusal.environment_variable_is_set else _SAVE_KEY_STEP - match refusal.source: - case ConfigFileSource() as source: + 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"1. 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}\n" - f"2. {set_key_step}" + 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 f"1. {set_key_step}" + return () case _: - assert_never(refusal.source) + 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, which encrypts stored\n" + f"credentials while {SALT_KEY_ENV_VAR} is not set. Replacing the key alone makes them unreadable, so also tell\n" + "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..e655d51b31e 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -119,6 +119,33 @@ def encrypt_value_helper(value: str, new_encryption_key: str | None = None): raise e +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) + + # 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) + + return decrypt_value(value=decoded_b64, 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. Both ciphers are authenticated, so a wrong key never passes.""" + if not value: + return None + try: + return _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 +156,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..452298516bd --- /dev/null +++ b/litellm/proxy/db/master_key_migration.py @@ -0,0 +1,248 @@ +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.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) -> tuple[JsonValue, int]: + 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) 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) 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 = current_schema()" + ) + ) + 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 + + +MigrationOutcome = NothingToMigrate | Migrated + + +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( + 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( + *, + 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 _: + assert_never(outcome) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6c734186faf..5eff79193f1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -344,11 +344,13 @@ from litellm.proxy.auth.login_throttle import ( ) from litellm.proxy.auth.master_key_boot_check import ( MASTER_KEY_ENV_VAR, + MIGRATE_FROM_MASTER_KEY_ENV_VAR, SALT_KEY_ENV_VAR, UNSAFE_PROXY_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, @@ -472,6 +474,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 ( @@ -486,6 +489,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_from_previous_master_key, +) from litellm.proxy.db.proxy_worker_heartbeat import ( PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS, ProxyWorkerHeartbeat, @@ -1131,6 +1138,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.db + + @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: global \ @@ -1224,14 +1239,17 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await initialize(**worker_config) enforce_master_key_boot_verdict( - 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(UNSAFE_PROXY_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, + 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(UNSAFE_PROXY_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, ) @@ -1245,6 +1263,16 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: user_api_key_cache=user_api_key_cache, ) + previous_master_key: Final = os.getenv(MIGRATE_FROM_MASTER_KEY_ENV_VAR) + if previous_master_key is not None and master_key is not None: + await migrate_from_previous_master_key( + previous_master_key=previous_master_key, + master_key=master_key, + salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, + database=None if prisma_client is None else prisma_client.db, + log=verbose_proxy_logger.warning, + ) + if prisma_client is not None: async def _run_pw_migration(): 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 index 2bb33276a0c..bf208bf235b 100644 --- a/tests/test_litellm/proxy/auth/test_master_key_boot_check.py +++ b/tests/test_litellm/proxy/auth/test_master_key_boot_check.py @@ -1,3 +1,4 @@ +import asyncio import re import shutil import subprocess @@ -10,6 +11,7 @@ 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, UNSAFE_PROXY_OVERRIDE_ENV_VAR, @@ -18,6 +20,7 @@ from litellm.proxy.auth.master_key_boot_check import ( EnvironmentSource, MasterKeyBootVerdict, SafeMasterKey, + StoredSecretsMigration, UnsafeMasterKeyAllowed, UnsafeMasterKeyError, UnsafeMasterKeyReason, @@ -26,6 +29,7 @@ from litellm.proxy.auth.master_key_boot_check import ( enforce_master_key_boot_verdict, master_key_boot_verdict, render_refusal, + with_stored_secrets_counted, ) @@ -125,38 +129,86 @@ def test_environment_is_the_source_when_yaml_does_not_set_a_master_key(): @pytest.mark.parametrize( - ("master_key", "salt_key_is_set", "database_is_configured", "needs_rotation"), + ("master_key", "salt_key_is_set", "database_is_configured", "migration"), [ - ("sk-1234", False, True, True), - ("sk-1234", True, True, False), - ("sk-1234", False, False, False), - (None, False, True, False), - ("", False, True, False), + ("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_rotation_is_only_needed_when_the_known_key_encrypts_a_database( - master_key: str | None, salt_key_is_set: bool, database_is_configured: bool, needs_rotation: bool +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.stored_credentials_need_rotation is needs_rotation + 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, - stored_credentials_need_rotation: bool = False, + migration: StoredSecretsMigration | None = None, ) -> UnsafeMasterKeyRefused: return UnsafeMasterKeyRefused( reason=reason, source=source, environment_variable_is_set=environment_variable_is_set, - stored_credentials_need_rotation=stored_credentials_need_rotation, + 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"))) @@ -208,28 +260,82 @@ def test_unset_key_refusal_says_nothing_supplied_one(): assert "Neither general_settings.master_key nor" in text -def test_rotation_guide_appears_only_when_needed(): - with_rotation = render_refusal(_refusal(stored_credentials_need_rotation=True)) - without_rotation = render_refusal(_refusal(stored_credentials_need_rotation=False)) +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 ROTATION_DOCS_URL in with_rotation - assert ROTATION_DOCS_URL not in without_rotation + 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 -@pytest.mark.parametrize("source", [EnvironmentSource(), ConfigFileSource(config_file_path="/app/config.yaml")]) -def test_refusal_never_tells_a_user_who_must_rotate_to_save_the_new_key_first( - source: ConfigFileSource | EnvironmentSource, -): - text = render_refusal(_refusal(source=source, stored_credentials_need_rotation=True)) +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 ".env" not in text - assert "os.environ/" not 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.parametrize("stored_credentials_need_rotation", [True, False]) -def test_override_hint_is_the_last_paragraph(stored_credentials_need_rotation: bool): - text = render_refusal(_refusal(stored_credentials_need_rotation=stored_credentials_need_rotation)) +@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 UNSAFE_PROXY_OVERRIDE_ENV_VAR in last_paragraph 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..fe9659f4fd7 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,33 @@ 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/db/test_master_key_migration.py b/tests/test_litellm/proxy/db/test_master_key_migration.py new file mode 100644 index 00000000000..31732adc666 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_master_key_migration.py @@ -0,0 +1,398 @@ +import json +import re +from collections.abc import Mapping, Sequence + +import pytest + +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, + NothingToMigrate, + count_values_encrypted_with, + describe_outcome, + migrate_from_previous_master_key, + 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 = current_schema()" 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_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"}] + + +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) == "" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index a670f467619..191fa862f6c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1633,7 +1633,14 @@ def _boot_with_general_settings(monkeypatch, tmp_path, general_settings): config_path = tmp_path / "config.yaml" config_path.write_text(yaml.dump({"general_settings": general_settings})) - for name in ("LITELLM_MASTER_KEY", "LITELLM_DANGEROUSLY_ALLOW_UNSAFE_PROXY", "WORKER_CONFIG", "DATABASE_URL"): + for name in ( + "LITELLM_MASTER_KEY", + "LITELLM_DANGEROUSLY_ALLOW_UNSAFE_PROXY", + "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" @@ -1649,7 +1656,7 @@ def _boot_with_general_settings(monkeypatch, tmp_path, 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_before_connecting_to_the_database( +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 @@ -1657,8 +1664,12 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_before_connecting_to_t 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()): @@ -1666,6 +1677,68 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_before_connecting_to_t 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] @pytest.mark.asyncio