mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
* refactor(proxy): expose public names for private proxy helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep original class names behind public aliases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): preserve internal callback filtering Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep _PROXY_ class names for managed files hooks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep old private names in package exports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): preserve recursive auth helper name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): align MCP limiter tests with server enforcement Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): match main's MCP limiter tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep old private names bound in importing modules Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): preserve compatibility imports through strict lint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): use exact pyright suppression in password helper test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): add reasons to compatibility import noqa comments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): update IN-list baseline for renamed helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
647 lines
25 KiB
Python
647 lines
25 KiB
Python
"""
|
|
Tests for the at-rest credential re-encryption migration engine.
|
|
|
|
The pure engine (classify / reencrypt / selective-dict) is tested directly; the
|
|
DB walkers are tested against an AsyncMock Prisma client. Live end-to-end
|
|
proof-of-fix (real proxy + DB) is performed separately on the repro server.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from types import SimpleNamespace
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from litellm._service_logger import ServiceTypes
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|
V2_GCM_PREFIX,
|
|
encrypt_value_helper,
|
|
)
|
|
from litellm.proxy.management_endpoints import credential_migration as cm
|
|
from tests.unit.proxy.db.fake_prisma_engine import engine_call
|
|
|
|
|
|
@pytest.fixture
|
|
def salt_key(monkeypatch):
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-migration-salt-1234")
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
return "sk-migration-salt-1234"
|
|
|
|
|
|
def _legacy_ct(value: str, monkeypatch) -> str:
|
|
"""Produce a legacy (nacl) ciphertext with the AES gate off."""
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
return encrypt_value_helper(value)
|
|
|
|
|
|
def _enable_aes(monkeypatch):
|
|
monkeypatch.setattr(
|
|
proxy_server, "general_settings", {"encryption_algorithm": "aes-256-gcm"}
|
|
)
|
|
|
|
|
|
def _empty_covered_tables(client):
|
|
"""Wire every rotation-covered table on `client` to return no rows.
|
|
|
|
Lets a `check_encryption` / scanner test isolate the location under test
|
|
without the other covered tables raising on an unconfigured mock.
|
|
"""
|
|
for _, db_attr, _, _ in cm._COVERED_TABLE_SPECS:
|
|
getattr(client.db, db_attr).find_many = AsyncMock(return_value=[])
|
|
|
|
|
|
# --------------------------- pure engine ---------------------------
|
|
|
|
|
|
def test_classify_value(salt_key, monkeypatch):
|
|
legacy = _legacy_ct("secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
migrated = encrypt_value_helper("secret")
|
|
|
|
assert cm.classify_value(legacy) == "legacy"
|
|
assert cm.classify_value(migrated) == "migrated"
|
|
assert cm.classify_value("just-plaintext") == "plaintext"
|
|
assert cm.classify_value("") == "plaintext"
|
|
assert cm.classify_value(123) == "not-a-string"
|
|
assert cm.classify_value(None) == "not-a-string"
|
|
|
|
|
|
def test_is_migrated(salt_key, monkeypatch):
|
|
_enable_aes(monkeypatch)
|
|
assert cm.is_migrated(encrypt_value_helper("x")) is True
|
|
assert cm.is_migrated("plaintext") is False
|
|
assert cm.is_migrated(5) is False
|
|
|
|
|
|
def test_reencrypt_value_legacy_to_v2(salt_key, monkeypatch):
|
|
legacy = _legacy_ct("secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
|
|
out = cm.reencrypt_value(legacy)
|
|
assert out != legacy
|
|
assert out.startswith(V2_GCM_PREFIX)
|
|
|
|
|
|
def test_reencrypt_value_is_idempotent(salt_key, monkeypatch):
|
|
_enable_aes(monkeypatch)
|
|
v2 = encrypt_value_helper("secret")
|
|
# Already v2 -> returned byte-for-byte unchanged (no re-wrap).
|
|
assert cm.reencrypt_value(v2) == v2
|
|
|
|
|
|
def test_reencrypt_value_preserves_non_string_and_empty(salt_key, monkeypatch):
|
|
_enable_aes(monkeypatch)
|
|
assert cm.reencrypt_value(42) == 42
|
|
assert cm.reencrypt_value("") == ""
|
|
assert cm.reencrypt_value(None) is None
|
|
|
|
|
|
def test_reencrypt_value_skips_undecryptable(salt_key, monkeypatch):
|
|
"""A value that does not decrypt (legacy plaintext or corrupt) is preserved."""
|
|
_enable_aes(monkeypatch)
|
|
plaintext = "not-actually-encrypted"
|
|
assert cm.reencrypt_value(plaintext) == plaintext
|
|
|
|
|
|
def test_reencrypt_selective_dict(salt_key, monkeypatch):
|
|
legacy_key = _legacy_ct("the-api-key", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
|
|
data = {"api_key": legacy_key, "base_url": "https://x", "integration_token": None}
|
|
out = cm.reencrypt_selective_dict(data, ["api_key", "integration_token"])
|
|
|
|
assert out["api_key"].startswith(V2_GCM_PREFIX)
|
|
assert out["base_url"] == "https://x" # untouched non-sensitive
|
|
assert out["integration_token"] is None # null skipped
|
|
|
|
|
|
# --------------------------- gate enforcement ---------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_migrate_requires_aes_gate(salt_key, monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "general_settings", {}) # gate off
|
|
with pytest.raises(RuntimeError, match="encryption_algorithm"):
|
|
await cm.migrate_encryption(
|
|
prisma_client=MagicMock(), user_api_key_dict=MagicMock()
|
|
)
|
|
|
|
|
|
# --------------------------- config-row walker ---------------------------
|
|
|
|
|
|
def _config_prisma(record):
|
|
"""Build an AsyncMock prisma client whose litellm_config returns `record`."""
|
|
client = MagicMock()
|
|
client.db.litellm_config.find_unique = AsyncMock(return_value=record)
|
|
client.db.litellm_config.update = engine_call()
|
|
return client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vantage_walker_migrates_legacy_field(salt_key, monkeypatch):
|
|
legacy_api_key = _legacy_ct("vantage-secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
record = SimpleNamespace(
|
|
param_value={
|
|
"api_key": legacy_api_key,
|
|
"integration_token": None,
|
|
"base_url": "https://api.vantage.sh",
|
|
}
|
|
)
|
|
client = _config_prisma(record)
|
|
|
|
report = await cm._migrate_config_settings_row(
|
|
client, "vantage_settings", cm._VANTAGE_SENSITIVE, dry_run=False
|
|
)
|
|
|
|
assert report.migrated == 1
|
|
assert report.legacy == 0 # migrated -> no longer residual legacy
|
|
client.db.litellm_config.update.assert_awaited_once()
|
|
written = json.loads(
|
|
client.db.litellm_config.update.call_args.kwargs["data"]["param_value"]
|
|
)
|
|
assert written["api_key"].startswith(V2_GCM_PREFIX)
|
|
assert written["base_url"] == "https://api.vantage.sh" # non-sensitive untouched
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vantage_walker_idempotent_no_write(salt_key, monkeypatch):
|
|
_enable_aes(monkeypatch)
|
|
record = SimpleNamespace(
|
|
param_value={"api_key": encrypt_value_helper("already-v2"), "base_url": "x"}
|
|
)
|
|
client = _config_prisma(record)
|
|
|
|
report = await cm._migrate_config_settings_row(
|
|
client, "vantage_settings", cm._VANTAGE_SENSITIVE, dry_run=False
|
|
)
|
|
|
|
assert report.already_v2 == 1
|
|
assert report.migrated == 0
|
|
client.db.litellm_config.update.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_config_walker_dry_run_does_not_write(salt_key, monkeypatch):
|
|
legacy_api_key = _legacy_ct("vantage-secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
record = SimpleNamespace(param_value={"api_key": legacy_api_key})
|
|
client = _config_prisma(record)
|
|
|
|
report = await cm._migrate_config_settings_row(
|
|
client, "vantage_settings", cm._VANTAGE_SENSITIVE, dry_run=True
|
|
)
|
|
|
|
# A dry run reports residual legacy only; nothing is migrated (no write), so
|
|
# `migrated` and `residual_legacy` are never contradictory in --check output.
|
|
assert report.legacy == 1
|
|
assert report.migrated == 0
|
|
client.db.litellm_config.update.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_config_walker_handles_missing_row(salt_key, monkeypatch):
|
|
_enable_aes(monkeypatch)
|
|
client = _config_prisma(None)
|
|
report = await cm._migrate_config_settings_row(
|
|
client, "cloudzero_settings", cm._CLOUDZERO_SENSITIVE, dry_run=False
|
|
)
|
|
assert report.scanned == 0
|
|
client.db.litellm_config.update.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sso_walker_real_run_migrates_and_clears_residual(salt_key, monkeypatch):
|
|
"""SSO real run: a migrated field is counted as migrated, not residual legacy."""
|
|
legacy = _legacy_ct("client-secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
record = SimpleNamespace(sso_settings={"client_secret": legacy, "client_id": "id"})
|
|
client = MagicMock()
|
|
client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=record)
|
|
client.db.litellm_ssoconfig.update = AsyncMock()
|
|
|
|
report = await cm._migrate_sso_config(client, dry_run=False)
|
|
|
|
assert report.migrated == 1
|
|
assert report.legacy == 0 # migrated -> no longer residual
|
|
client.db.litellm_ssoconfig.update.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sso_walker_dry_run_reports_residual_not_migrated(salt_key, monkeypatch):
|
|
"""SSO dry run: residual legacy only; migrated stays 0 (never contradictory)."""
|
|
legacy = _legacy_ct("client-secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
record = SimpleNamespace(sso_settings={"client_secret": legacy})
|
|
client = MagicMock()
|
|
client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=record)
|
|
client.db.litellm_ssoconfig.update = AsyncMock()
|
|
|
|
report = await cm._migrate_sso_config(client, dry_run=True)
|
|
|
|
assert report.legacy == 1
|
|
assert report.migrated == 0
|
|
client.db.litellm_ssoconfig.update.assert_not_awaited()
|
|
|
|
|
|
# --------------------------- --check scanner ---------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_reports_residual_legacy(salt_key, monkeypatch):
|
|
legacy_api_key = _legacy_ct("vantage-secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
|
|
client = MagicMock()
|
|
# Net-new walker tables: empty team / token / sso, one legacy vantage field.
|
|
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.update = AsyncMock()
|
|
_empty_covered_tables(client)
|
|
|
|
def _find_unique(where):
|
|
if where.get("param_name") == "vantage_settings":
|
|
return SimpleNamespace(param_value={"api_key": legacy_api_key})
|
|
return None
|
|
|
|
client.db.litellm_config.find_unique = AsyncMock(side_effect=_find_unique)
|
|
|
|
report = await cm.check_encryption(client)
|
|
|
|
assert report.residual_legacy == 1
|
|
client.db.litellm_config.update.assert_not_awaited() # read-only
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_reports_zero_after_migration(salt_key, monkeypatch):
|
|
_enable_aes(monkeypatch)
|
|
client = MagicMock()
|
|
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.update = AsyncMock()
|
|
_empty_covered_tables(client)
|
|
|
|
def _find_unique(where):
|
|
if where.get("param_name") == "vantage_settings":
|
|
return SimpleNamespace(
|
|
param_value={"api_key": encrypt_value_helper("already-v2")}
|
|
)
|
|
return None
|
|
|
|
client.db.litellm_config.find_unique = AsyncMock(side_effect=_find_unique)
|
|
|
|
report = await cm.check_encryption(client)
|
|
assert report.residual_legacy == 0
|
|
|
|
|
|
# --------------------------- callback_vars walker ---------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_callback_vars_walker_migrates_team_metadata(salt_key, monkeypatch):
|
|
"""A team row with a legacy-encrypted callback var is rewritten to v2."""
|
|
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
|
|
|
|
# Legacy-encrypt a callback var via the real callback path (gate off).
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
legacy_meta = encrypt_callback_vars(
|
|
{"logging": [{"callback_vars": {"gcs_path_service_account": "sa-secret"}}]}
|
|
)
|
|
_enable_aes(monkeypatch)
|
|
|
|
team_row = SimpleNamespace(team_id="team-1", metadata=legacy_meta)
|
|
client = MagicMock()
|
|
client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_row])
|
|
client.db.litellm_teamtable.update = AsyncMock()
|
|
|
|
report = await cm._migrate_callback_vars_table(client, "team", dry_run=False)
|
|
|
|
assert report.migrated == 1
|
|
assert report.scanned == 1 # one field examined, not "post-v2" count
|
|
client.db.litellm_teamtable.update.assert_awaited_once()
|
|
written = json.loads(
|
|
client.db.litellm_teamtable.update.call_args.kwargs["data"]["metadata"]
|
|
)
|
|
inner = written["logging"][0]["callback_vars"]["gcs_path_service_account"]
|
|
assert "v2:gcm:" in inner
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_callback_vars_walker_dry_run_reports_legacy(salt_key, monkeypatch):
|
|
"""In --check (dry-run) mode, a legacy callback var counts as residual legacy."""
|
|
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
|
|
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
legacy_meta = encrypt_callback_vars(
|
|
{"logging": [{"callback_vars": {"gcs_path_service_account": "sa-secret"}}]}
|
|
)
|
|
_enable_aes(monkeypatch)
|
|
|
|
team_row = SimpleNamespace(team_id="team-1", metadata=legacy_meta)
|
|
client = MagicMock()
|
|
client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_row])
|
|
client.db.litellm_teamtable.update = AsyncMock()
|
|
|
|
report = await cm._migrate_callback_vars_table(client, "team", dry_run=True)
|
|
|
|
assert report.scanned == 1
|
|
assert report.legacy == 1 # would-migrate -> residual legacy in attestation
|
|
assert report.migrated == 0
|
|
client.db.litellm_teamtable.update.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_callback_vars_walker_migrates_callback_settings_shape(
|
|
salt_key, monkeypatch
|
|
):
|
|
"""Regression: credentials under ``metadata.callback_settings.callback_vars``
|
|
with no top-level ``logging`` key must be migrated, not skipped.
|
|
|
|
The walker previously early-continued on ``"logging" not in metadata``, so
|
|
this credential shape (which ``encrypt_callback_vars`` does encrypt) was left
|
|
in legacy format at rest while the migration still reported success.
|
|
"""
|
|
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
|
|
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
legacy_meta = encrypt_callback_vars(
|
|
{
|
|
"callback_settings": {
|
|
"callback_vars": {"gcs_path_service_account": "sa-secret"}
|
|
}
|
|
}
|
|
)
|
|
_enable_aes(monkeypatch)
|
|
assert "logging" not in legacy_meta # the shape that used to be skipped
|
|
|
|
team_row = SimpleNamespace(team_id="team-1", metadata=legacy_meta)
|
|
client = MagicMock()
|
|
client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_row])
|
|
client.db.litellm_teamtable.update = AsyncMock()
|
|
|
|
report = await cm._migrate_callback_vars_table(client, "team", dry_run=False)
|
|
|
|
assert report.migrated == 1
|
|
assert report.scanned == 1
|
|
client.db.litellm_teamtable.update.assert_awaited_once()
|
|
written = json.loads(
|
|
client.db.litellm_teamtable.update.call_args.kwargs["data"]["metadata"]
|
|
)
|
|
inner = written["callback_settings"]["callback_vars"]["gcs_path_service_account"]
|
|
assert "v2:gcm:" in inner
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_reports_callback_var_legacy_with_gate_off(salt_key, monkeypatch):
|
|
"""check_encryption must report residual legacy callback vars even when the
|
|
AES gate is OFF.
|
|
|
|
Detection is decrypt-based, not a re-encrypt delta, so it does not depend on
|
|
the write gate. A heuristic that re-encrypts and counts new v2 values would
|
|
read zero here (gate off -> no v2 produced) and emit a false-clean
|
|
attestation -- exactly the compliance trap this guards against.
|
|
"""
|
|
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
|
|
|
|
# Legacy-encrypt a callback var, and leave the gate OFF for the check itself.
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
legacy_meta = encrypt_callback_vars(
|
|
{"logging": [{"callback_vars": {"gcs_path_service_account": "sa-secret"}}]}
|
|
)
|
|
team_row = SimpleNamespace(team_id="team-1", metadata=legacy_meta)
|
|
|
|
client = MagicMock()
|
|
_empty_covered_tables(client)
|
|
client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_row])
|
|
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)
|
|
client.db.litellm_teamtable.update = AsyncMock()
|
|
|
|
report = await cm.check_encryption(client)
|
|
|
|
assert report.residual_legacy == 1
|
|
assert report.as_dict()["locations"]["team.callback_vars"]["legacy"] == 1
|
|
client.db.litellm_teamtable.update.assert_not_awaited() # read-only
|
|
|
|
|
|
# --------------------------- covered-tables scanner ---------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scan_covered_tables_classifies_legacy_and_v2(salt_key, monkeypatch):
|
|
"""The read-only scanner classifies the model and credentials tables."""
|
|
legacy = _legacy_ct("model-secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
v2 = encrypt_value_helper("cred-secret")
|
|
|
|
client = MagicMock()
|
|
_empty_covered_tables(client)
|
|
client.db.litellm_proxymodeltable.find_many = AsyncMock(
|
|
return_value=[
|
|
SimpleNamespace(litellm_params={"api_key": legacy, "model": "gpt-4"})
|
|
]
|
|
)
|
|
client.db.litellm_credentialstable.find_many = AsyncMock(
|
|
return_value=[SimpleNamespace(credential_values={"api_key": v2})]
|
|
)
|
|
client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
|
|
|
by_loc = {r.location: r for r in await cm._scan_covered_tables(client)}
|
|
|
|
assert by_loc["model_table"].legacy == 1
|
|
assert by_loc["model_table"].plaintext == 1 # "gpt-4" model name, not ciphertext
|
|
assert by_loc["credentials"].already_v2 == 1
|
|
assert by_loc["credentials"].legacy == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scan_covered_tables_classifies_search_tool_params(salt_key, monkeypatch):
|
|
legacy = _legacy_ct("tvly-legacy", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
v2 = encrypt_value_helper("tvly-migrated")
|
|
|
|
client = MagicMock()
|
|
_empty_covered_tables(client)
|
|
client.db.litellm_searchtoolstable.find_many = AsyncMock(
|
|
return_value=[
|
|
SimpleNamespace(litellm_params={"api_key": legacy, "timeout": 30}),
|
|
SimpleNamespace(litellm_params={"api_key": v2, "search_provider": "tavily"}),
|
|
]
|
|
)
|
|
client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
|
|
|
by_loc = {r.location: r for r in await cm._scan_covered_tables(client)}
|
|
|
|
assert (by_loc["search_tools"].legacy, by_loc["search_tools"].already_v2) == (1, 1)
|
|
assert by_loc["search_tools"].plaintext == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("column", ("static_headers", "env"))
|
|
@pytest.mark.parametrize("algorithm", ("xsalsa20-poly1305", "aes-256-gcm"))
|
|
@pytest.mark.parametrize("as_json", (False, True))
|
|
@pytest.mark.parametrize(
|
|
"case", ("legacy", "encrypted", "wrong-key", "corrupt", "invalid-shape", "invalid-scalar", "empty", "null")
|
|
)
|
|
async def test_check_classifies_mcp_secret_maps(
|
|
salt_key: str,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
column: str,
|
|
algorithm: str,
|
|
as_json: bool,
|
|
case: str,
|
|
) -> None:
|
|
monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": algorithm})
|
|
plaintext: Final = {"Authorization": "v2:gcm:operator-text", "CUSTOM": "litellm_enc::literal\n café "}
|
|
ciphertext: Final = encrypt_value_helper(json.dumps(plaintext))
|
|
cases: Final[dict[str, object]] = {
|
|
"legacy": plaintext,
|
|
"encrypted": ciphertext,
|
|
"wrong-key": encrypt_value_helper(json.dumps(plaintext), new_encryption_key="different-map-salt"),
|
|
"corrupt": ciphertext[:-4] + "AAAA",
|
|
"invalid-shape": encrypt_value_helper(json.dumps({"Authorization": 42})),
|
|
"invalid-scalar": "null",
|
|
"empty": {},
|
|
"null": None,
|
|
}
|
|
value: Final = json.dumps(cases[case]) if as_json and case != "null" else cases[case]
|
|
row: Final = SimpleNamespace(**{column: value})
|
|
client: Final = MagicMock()
|
|
_empty_covered_tables(client)
|
|
client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row])
|
|
client.db.litellm_mcpservertable.update = AsyncMock()
|
|
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)
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
|
|
report: Final = await cm.check_encryption(client)
|
|
expected_legacy: Final = int(case == "legacy" or (case == "encrypted" and algorithm == "xsalsa20-poly1305"))
|
|
expected_v2: Final = int(case == "encrypted" and algorithm == "aes-256-gcm")
|
|
expected_invalid: Final = int(case in ("wrong-key", "corrupt", "invalid-shape", "invalid-scalar"))
|
|
|
|
assert report.as_dict()["locations"]["mcp_server"] == {
|
|
"scanned": int(case not in ("empty", "null")),
|
|
"migrated": 0,
|
|
"already_v2": expected_v2,
|
|
"plaintext": 0,
|
|
"undecryptable": expected_invalid,
|
|
"legacy": expected_legacy,
|
|
}
|
|
assert report.residual_legacy == expected_legacy
|
|
assert report.total_undecryptable == expected_invalid
|
|
assert getattr(row, column) == value
|
|
client.db.litellm_mcpservertable.update.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_counts_covered_table_residual(salt_key, monkeypatch):
|
|
"""check_encryption now scans the rotation-covered tables (model table here),
|
|
so a legacy value there counts toward residual_legacy (the P1 attestation gap).
|
|
"""
|
|
legacy = _legacy_ct("model-secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
|
|
client = 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)
|
|
client.db.litellm_config.update = AsyncMock()
|
|
client.db.litellm_proxymodeltable.find_many = AsyncMock(
|
|
return_value=[SimpleNamespace(litellm_params={"api_key": legacy})]
|
|
)
|
|
|
|
report = await cm.check_encryption(client)
|
|
|
|
assert report.residual_legacy == 1
|
|
assert report.as_dict()["locations"]["model_table"]["legacy"] == 1
|
|
client.db.litellm_config.update.assert_not_awaited() # read-only
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_migrate_covered_tables_reports_real_counts(salt_key, monkeypatch):
|
|
"""_migrate_covered_tables derives real per-table counts from pre/post scans,
|
|
instead of the always-zero report Greptile flagged (P1).
|
|
"""
|
|
legacy = _legacy_ct("model-secret", monkeypatch)
|
|
_enable_aes(monkeypatch)
|
|
v2 = encrypt_value_helper("model-secret")
|
|
|
|
row = SimpleNamespace(litellm_params={"api_key": legacy})
|
|
client = MagicMock()
|
|
_empty_covered_tables(client)
|
|
client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[row])
|
|
client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
|
|
|
async def fake_rotate(**kwargs):
|
|
# Stand in for rotate_master_key: re-encrypt the model api_key in place.
|
|
row.litellm_params["api_key"] = v2
|
|
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.management_endpoints.key_management_endpoints.rotate_master_key",
|
|
fake_rotate,
|
|
)
|
|
|
|
by_loc = {
|
|
r.location: r for r in await cm._migrate_covered_tables(client, MagicMock())
|
|
}
|
|
|
|
assert by_loc["model_table"].migrated == 1 # was legacy pre, v2 post
|
|
assert by_loc["model_table"].legacy == 0 # residual zero after rotation
|
|
assert by_loc["model_table"].already_v2 == 1
|
|
|
|
|
|
def _db_service_hooks() -> tuple[AsyncMock, MagicMock]:
|
|
success: Final = AsyncMock()
|
|
service_logging: Final = MagicMock(async_service_success_hook=success, async_service_failure_hook=AsyncMock())
|
|
return success, MagicMock(service_logging_obj=service_logging)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_config_walker_write_emits_a_postgres_update_event_for_litellm_config(salt_key, monkeypatch):
|
|
_enable_aes(monkeypatch)
|
|
client = _config_prisma(SimpleNamespace(param_value={"api_key": _legacy_ct("vantage-secret", monkeypatch)}))
|
|
success, proxy_logging = _db_service_hooks()
|
|
monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging)
|
|
|
|
await cm._migrate_config_settings_row(client, "vantage_settings", cm._VANTAGE_SENSITIVE, dry_run=False)
|
|
await asyncio.sleep(0)
|
|
|
|
event: Final = success.await_args.kwargs
|
|
assert (event["service"], event["call_type"], event["event_metadata"]) == (
|
|
ServiceTypes.DB,
|
|
"migrate_config_credentials",
|
|
{"table_name": "LiteLLM_Config"},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sso_walker_write_emits_a_postgres_update_event_for_litellm_ssoconfig(salt_key, monkeypatch):
|
|
_enable_aes(monkeypatch)
|
|
client = MagicMock()
|
|
client.db.litellm_ssoconfig.find_unique = AsyncMock(
|
|
return_value=SimpleNamespace(sso_settings={"client_secret": _legacy_ct("client-secret", monkeypatch)})
|
|
)
|
|
client.db.litellm_ssoconfig.update = engine_call()
|
|
success, proxy_logging = _db_service_hooks()
|
|
monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging)
|
|
|
|
await cm._migrate_sso_config(client, dry_run=False)
|
|
await asyncio.sleep(0)
|
|
|
|
event: Final = success.await_args.kwargs
|
|
assert (event["service"], event["call_type"], event["event_metadata"]) == (
|
|
ServiceTypes.DB,
|
|
"migrate_sso_credentials",
|
|
{"table_name": "LiteLLM_SSOConfig"},
|
|
)
|