diff --git a/litellm/proxy/management_endpoints/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py index 2e98364c7db..db6e518a663 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -462,6 +462,8 @@ _COVERED_TABLE_SPECS: Final = [ ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()), ("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)), ("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)), + ("mcp_oauth_client", "litellm_mcpserveroauthclient", ("credentials",), ()), + ("sso_identity_assertion", "litellm_ssoidentityassertion", (), ("assertion_b64",)), ] @@ -493,14 +495,27 @@ def _classify_into_report(report: LocationReport, value: str) -> None: report.count(classify_value(value, key="scan")) +def _secret_map_ciphertext(value: object) -> object: + """The stored ciphertext behind an MCP secret map: the JSON quoted string is unwrapped, a plain map or + anything that is not JSON is returned as is for ``decode_secret_map`` to judge.""" + if not isinstance(value, str) or not value.lstrip().startswith('"'): + return value + try: + return json.loads(value) + except ValueError: + return value + + def _classify_secret_map(value: object, key: str) -> ValueClass | None: + ciphertext: Final = _secret_map_ciphertext(value) + if isinstance(ciphertext, str) and not ciphertext.lstrip().startswith("{"): + require_legacy_reader_for((ciphertext,), "scan stored encryption") try: decoded: Final = decode_secret_map(value, key=key) except SecretMapDecodeError: return "undecryptable" if not decoded: return None - ciphertext: Final = json.loads(value) if isinstance(value, str) and value.lstrip().startswith('"') else value return "migrated" if is_migrated(ciphertext) else "legacy" diff --git a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py index eb00462141e..24eb3cdab78 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py +++ b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py @@ -17,6 +17,7 @@ import pytest from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import ( _V3_GCM_PREFIX, + encrypt_secret_map, encrypt_value_helper, ) from litellm.proxy.management_endpoints import credential_migration as cm @@ -157,6 +158,41 @@ async def test_migrate_and_check_refuse_without_pynacl_instead_of_miscounting_le client.db.litellm_config.update.assert_not_awaited() +@pytest.mark.parametrize( + "db_attr, build_row", + [ + ("litellm_mcpservertable", lambda ct: SimpleNamespace(server_id="s1", static_headers=json.dumps(ct))), + ("litellm_mcpservertable", lambda ct: SimpleNamespace(server_id="s1", env=ct)), + ("litellm_mcpserveroauthclient", lambda ct: SimpleNamespace(server_id="s1", credentials={"client_secret": ct})), + ("litellm_ssoidentityassertion", lambda ct: SimpleNamespace(user_id="u1", assertion_b64=ct)), + ], +) +@pytest.mark.asyncio +async def test_check_refuses_without_pynacl_for_every_rotation_rewritten_location( + db_attr, build_row, salt_key, monkeypatch +): + monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": "xsalsa20-poly1305"}) + legacy_map: Final = json.loads(encrypt_secret_map({"Authorization": "Bearer legacy"})) + legacy_value: Final = encrypt_value_helper("legacy-secret") + _enable_aes(monkeypatch) + client: Final = MagicMock() + _empty_covered_tables(client) + client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + client.db.litellm_config.find_unique = AsyncMock(return_value=None) + ciphertext: Final = legacy_map if db_attr == "litellm_mcpservertable" else legacy_value + getattr(client.db, db_attr).find_many = AsyncMock(return_value=[build_row(ciphertext)]) + + with_pynacl: Final = await cm.check_encryption(prisma_client=client) + assert with_pynacl.residual_legacy == 1, with_pynacl.to_dict() + + monkeypatch.setitem(sys.modules, "nacl", None) + monkeypatch.setitem(sys.modules, "nacl.secret", None) + with pytest.raises(RuntimeError, match="legacy-encryption"): + await cm.check_encryption(prisma_client=client) + + # --------------------------- config-row walker --------------------------- diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 46602c10888..882e38343a6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -9435,9 +9435,16 @@ async def test_rotate_master_key_reencrypts_model_params_in_place( ), "api_key must be stored re-encrypted under the new master key, not in plaintext" -async def test_rotate_master_key_refuses_without_pynacl_before_touching_any_row(monkeypatch): +@pytest.mark.parametrize( + "legacy_table, build_row", + [ + ("litellm_proxymodeltable", lambda ct: SimpleNamespace(litellm_params={"api_key": ct})), + ("litellm_mcpserveroauthclient", lambda ct: SimpleNamespace(server_id="s1", credentials={"client_secret": ct})), + ("litellm_ssoidentityassertion", lambda ct: SimpleNamespace(user_id="u1", assertion_b64=ct)), + ], +) +async def test_rotate_master_key_refuses_without_pynacl_before_touching_any_row(legacy_table, build_row, monkeypatch): import sys - from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock from fastapi import HTTPException @@ -9451,23 +9458,27 @@ async def test_rotate_master_key_refuses_without_pynacl_before_touching_any_row( monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-rotate") monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": "xsalsa20-poly1305"}) - legacy_row = SimpleNamespace(litellm_params={"api_key": encrypt_value_helper("legacy-secret")}) + legacy_row = build_row(encrypt_value_helper("legacy-secret")) monkeypatch.setattr(proxy_server, "general_settings", {}) monkeypatch.setitem(sys.modules, "nacl", None) monkeypatch.setitem(sys.modules, "nacl.secret", None) mock_prisma_client = AsyncMock() mock_prisma_client.db = MagicMock() - mock_prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[legacy_row]) - mock_prisma_client.db.litellm_proxymodeltable.update_many = AsyncMock() mock_prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) - empty_tables = ( + covered_tables = ( + "litellm_proxymodeltable", "litellm_credentialstable", "litellm_mcpservertable", "litellm_mcpusercredentials", "litellm_mcpuserenvvars", + "litellm_mcpserveroauthclient", + "litellm_ssoidentityassertion", ) - for table in empty_tables: - getattr(mock_prisma_client.db, table).find_many = AsyncMock(return_value=[]) + for table in covered_tables: + rows = [legacy_row] if table == legacy_table else [] + getattr(mock_prisma_client.db, table).find_many = AsyncMock(return_value=rows) + getattr(mock_prisma_client.db, table).update = AsyncMock() + getattr(mock_prisma_client.db, table).update_many = AsyncMock() with pytest.raises(HTTPException) as exc_info: await _rotate_master_key( @@ -9479,7 +9490,9 @@ async def test_rotate_master_key_refuses_without_pynacl_before_touching_any_row( assert exc_info.value.status_code == 400 assert "legacy-encryption" in exc_info.value.detail["error"] - mock_prisma_client.db.litellm_proxymodeltable.update_many.assert_not_awaited() + for table in covered_tables: + getattr(mock_prisma_client.db, table).update.assert_not_awaited() + getattr(mock_prisma_client.db, table).update_many.assert_not_awaited() mock_prisma_client.db.tx.assert_not_called()