Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_endpoints/credential_migration.py: 63%
319 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2At-rest credential re-encryption migration.
4Switches every encrypted-at-rest value from the legacy XSalsa20-Poly1305 (nacl)
5format to the versioned AES-256-GCM (``v2:gcm:``) format produced by
6``encrypt_decrypt_utils`` when ``general_settings.encryption_algorithm`` is set to
7``aes-256-gcm``.
9Design properties (see case 2026-06-24 fix plan):
11* **Same key, new algorithm.** The migration does not change the encryption key;
12 it re-encrypts existing ciphertext under the same derived key but in the new
13 AES format. This is achieved by decrypting with the format-detecting reader and
14 re-encrypting through ``encrypt_value_helper`` with the AES gate enabled.
15* **Idempotent.** A value already carrying the ``v2:gcm:`` prefix is recognised
16 and left untouched, so re-running the migration is a no-op on migrated rows.
17* **Resumable.** Walkers commit per row (or per small table), so an interrupted
18 run leaves a clean mixed state that a re-run completes.
19* **Skip-on-undecryptable.** A value that cannot be decrypted is never
20 overwritten — corrupt rows are preserved and reported, never destroyed.
21* **Attestable.** :func:`check_encryption` is a read-only scan that classifies
22 every value as ``migrated`` / ``legacy`` / ``plaintext`` / ``undecryptable``.
23 A residual ``legacy == 0`` is the compliance attestation.
25Coverage. The covered tables (model table, credentials table, MCP credential/env
26tables, config ``environment_variables``) already have a re-encryption path in
27``_rotate_master_key``; this module delegates to it in *same-key* mode and adds
28walkers for the locations that had no rotation path: team / verification-token
29``callback_vars`` metadata, the ``vantage_settings`` / ``cloudzero_settings``
30config rows, and the SSO config table.
31"""
33import json
34from dataclasses import dataclass, field
35from typing import TYPE_CHECKING, Final, Literal, cast
37from litellm._logging import verbose_proxy_logger
39if TYPE_CHECKING: 39 ↛ 40line 39 didn't jump to line 40 because the condition on line 39 was never true
40 from litellm.proxy._types import UserAPIKeyAuth
41 from litellm.proxy.utils import PrismaClient
42from litellm.proxy.common_utils.encrypt_decrypt_utils import (
43 _ALGO_AES_GCM,
44 _ENCRYPTION_ALGORITHM_SETTING,
45 _V2_GCM_PREFIX,
46 SecretMapDecodeError,
47 _get_salt_key,
48 decode_secret_map,
49 decrypt_value_helper,
50 encrypt_value_helper,
51)
53ValueClass = Literal["migrated", "legacy", "plaintext", "undecryptable", "not-a-string"]
56@dataclass
57class LocationReport:
58 """Per-location counters for one migration / check pass."""
60 location: str
61 scanned: int = 0
62 migrated: int = 0 # values rewritten to v2 this run
63 already_v2: int = 0 # values already migrated (skipped)
64 plaintext: int = 0 # legacy-plaintext values (no ciphertext to migrate)
65 undecryptable: int = 0 # could not decrypt — preserved, not overwritten
67 # Used by --check (read-only classification):
68 legacy: int = 0 # nacl ciphertext still awaiting migration
70 def count(self, classification: ValueClass | None) -> None:
71 if classification is None:
72 return
73 self.scanned += 1
74 match classification:
75 case "migrated": 75 ↛ 76line 75 didn't jump to line 76 because the pattern on line 75 never matched
76 self.already_v2 += 1
77 case "legacy":
78 self.legacy += 1
79 case "undecryptable": 79 ↛ 80line 79 didn't jump to line 80 because the pattern on line 79 never matched
80 self.undecryptable += 1
81 case _:
82 self.plaintext += 1
84 def as_dict(self) -> dict[str, int]:
85 return {
86 "scanned": self.scanned,
87 "migrated": self.migrated,
88 "already_v2": self.already_v2,
89 "plaintext": self.plaintext,
90 "undecryptable": self.undecryptable,
91 "legacy": self.legacy,
92 }
95@dataclass
96class MigrationReport:
97 """Aggregate report across all locations."""
99 locations: list[LocationReport] = field(default_factory=list)
101 def add(self, report: LocationReport) -> None:
102 self.locations.append(report)
104 @property
105 def residual_legacy(self) -> int:
106 """Total legacy ciphertext still un-migrated (the TRO attestation number)."""
107 return sum(loc.legacy for loc in self.locations)
109 @property
110 def total_undecryptable(self) -> int:
111 return sum(loc.undecryptable for loc in self.locations)
113 def as_dict(self) -> dict[str, object]:
114 return {
115 "residual_legacy": self.residual_legacy,
116 "total_undecryptable": self.total_undecryptable,
117 "locations": {loc.location: loc.as_dict() for loc in self.locations},
118 }
121# ---------------------------------------------------------------------------
122# Pure engine — no DB I/O, fully unit-testable.
123# ---------------------------------------------------------------------------
126def is_migrated(value: object) -> bool:
127 """True if ``value`` is already an AES-256-GCM (``v2:gcm:``) ciphertext."""
128 return isinstance(value, str) and value.startswith(_V2_GCM_PREFIX)
131def classify_value(value: object, key: str = "scan") -> ValueClass:
132 """Classify a stored value for the residual scanner.
134 * ``not-a-string`` — not a string (numbers/bools/None left as-is on disk).
135 * ``migrated`` — carries the ``v2:gcm:`` prefix.
136 * ``legacy`` — decrypts under the legacy nacl reader (still needs migrating).
137 * ``plaintext`` — a non-empty string that does not decrypt and is not v2;
138 treated as legacy plaintext (nothing to migrate).
139 * ``undecryptable`` — reserved for callers that already know a value is
140 ciphertext but cannot decrypt it; ``classify_value`` itself cannot tell a
141 corrupt ciphertext from plaintext, so it returns ``plaintext`` for both.
142 """
143 if not isinstance(value, str): 143 ↛ 144line 143 didn't jump to line 144 because the condition on line 143 was never true
144 return "not-a-string"
145 if value == "":
146 return "plaintext"
147 if value.startswith(_V2_GCM_PREFIX): 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true
148 return "migrated"
149 decrypted: Final = decrypt_value_helper(value=value, key=key, exception_type="debug", return_original_value=False)
150 if decrypted is None:
151 # Did not decrypt under nacl and has no v2 marker: legacy plaintext.
152 return "plaintext"
153 return "legacy"
156def reencrypt_value(value: object, key: str = "migrate") -> object:
157 """Re-encrypt a single stored string into the configured (AES) format.
159 Returns the value unchanged if it is not a string, is already ``v2:``, or
160 cannot be decrypted (skip-on-undecryptable). Otherwise decrypts under the
161 format-detecting reader and re-encrypts through ``encrypt_value_helper``
162 (which writes AES when the gate is on).
163 """
164 if not isinstance(value, str) or value == "":
165 return value
166 if value.startswith(_V2_GCM_PREFIX):
167 return value # idempotent: already migrated
168 decrypted: Final = decrypt_value_helper(value=value, key=key, exception_type="debug", return_original_value=False)
169 if decrypted is None:
170 # Either legacy plaintext (no ciphertext to migrate) or corrupt. Either
171 # way, do not overwrite — preserve the value as stored.
172 return value
173 return encrypt_value_helper(decrypted)
176def reencrypt_selective_dict(data: dict[str, object], sensitive_keys: list[str]) -> dict[str, object]:
177 """Return a copy of ``data`` with only ``sensitive_keys`` re-encrypted.
179 Non-sensitive fields (e.g. ``base_url``, ``connection_id``) are left as-is.
180 Null/missing fields are skipped.
181 """
182 out: Final = dict(data)
183 for k in sensitive_keys:
184 v = out.get(k)
185 if v is None:
186 continue
187 out[k] = reencrypt_value(v, key=k)
188 return out
191def _assert_aes_gate_enabled() -> None:
192 """Fail fast if the AES algorithm gate is not enabled.
194 Running the migration with the gate off would decrypt then re-encrypt right
195 back into the legacy format — a no-op that silently fails the migration.
196 """
197 from litellm.proxy.proxy_server import general_settings
199 algo: Final = general_settings.get(_ENCRYPTION_ALGORITHM_SETTING)
200 if not (isinstance(algo, str) and algo.lower() == _ALGO_AES_GCM): 200 ↛ exitline 200 didn't return from function '_assert_aes_gate_enabled' because the condition on line 200 was always true
201 raise RuntimeError(
202 "Encryption migration requires general_settings.encryption_algorithm: "
203 f"'{_ALGO_AES_GCM}'. Current value: {algo!r}. Set it before migrating "
204 "so re-encrypted values are written in the AES-256-GCM format."
205 )
208# ---------------------------------------------------------------------------
209# Walkers for the locations with no pre-existing rotation path.
210# Each walker delegates the structural transform to the existing, tested helper
211# for that table and only adds the per-row re-encrypt + commit + counters.
212# ---------------------------------------------------------------------------
215async def _migrate_config_settings_row(
216 prisma_client: object,
217 param_name: str,
218 sensitive_fields: list[str],
219 dry_run: bool,
220) -> LocationReport:
221 """Migrate a single ``LiteLLM_Config`` row whose ``param_value`` is a JSON
222 dict with selected sensitive fields (vantage_settings / cloudzero_settings).
223 """
224 report: Final = LocationReport(location=param_name)
225 record: Final = await prisma_client.db.litellm_config.find_unique(where={"param_name": param_name})
226 if record is None or record.param_value is None:
227 return report
229 settings = record.param_value
230 if isinstance(settings, str): 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true
231 settings = json.loads(settings)
232 if not isinstance(settings, dict): 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true
233 return report
235 changed = False
236 for fld in sensitive_fields:
237 v = settings.get(fld)
238 if v is None: 238 ↛ 239line 238 didn't jump to line 239 because the condition on line 238 was never true
239 continue
240 report.scanned += 1
241 cls = classify_value(v, key=fld)
242 if cls == "migrated": 242 ↛ 243line 242 didn't jump to line 243 because the condition on line 242 was never true
243 report.already_v2 += 1
244 continue
245 if cls == "legacy": 245 ↛ 261line 245 didn't jump to line 261 because the condition on line 245 was always true
246 if dry_run: 246 ↛ 251line 246 didn't jump to line 251 because the condition on line 246 was always true
247 # Residual: would migrate, but a dry run writes nothing, so it
248 # stays legacy for the attestation (never counted as migrated).
249 report.legacy += 1
250 continue
251 new_v = reencrypt_value(v, key=fld)
252 if new_v != v:
253 settings[fld] = new_v
254 report.migrated += 1
255 changed = True
256 else:
257 # Defensive: a legacy value that did not re-encrypt is still
258 # residual, not migrated.
259 report.legacy += 1
260 else: # plaintext / not-a-string — nothing to migrate
261 report.plaintext += 1
263 if changed and not dry_run: 263 ↛ 264line 263 didn't jump to line 264 because the condition on line 263 was never true
264 await prisma_client.db.litellm_config.update(
265 where={"param_name": param_name},
266 data={"param_value": json.dumps(settings)},
267 )
268 return report
271async def _migrate_sso_config(prisma_client: object, dry_run: bool) -> LocationReport:
272 """Migrate the ``LiteLLM_SSOConfig`` row. All non-null fields are encrypted
273 (via the same ``_encrypt_env_variables`` path used on save), so we re-encrypt
274 every present string field.
275 """
276 report: Final = LocationReport(location="sso_config")
277 record: Final = await prisma_client.db.litellm_ssoconfig.find_unique(where={"id": "sso_config"})
278 if record is None or record.sso_settings is None:
279 return report
281 settings = record.sso_settings
282 if isinstance(settings, str): 282 ↛ 283line 282 didn't jump to line 283 because the condition on line 282 was never true
283 settings = json.loads(settings)
284 if not isinstance(settings, dict): 284 ↛ 285line 284 didn't jump to line 285 because the condition on line 284 was never true
285 return report
287 new_settings: Final = dict(settings)
288 changed = False
289 for fld, v in settings.items():
290 if not isinstance(v, str) or v == "":
291 continue
292 report.scanned += 1
293 cls = classify_value(v, key=fld)
294 if cls == "migrated": 294 ↛ 295line 294 didn't jump to line 295 because the condition on line 294 was never true
295 report.already_v2 += 1
296 continue
297 if cls == "legacy": 297 ↛ 313line 297 didn't jump to line 313 because the condition on line 297 was always true
298 if dry_run: 298 ↛ 303line 298 didn't jump to line 303 because the condition on line 298 was always true
299 # Residual: would migrate, but a dry run writes nothing, so it
300 # stays legacy for the attestation (never counted as migrated).
301 report.legacy += 1
302 continue
303 new_v = reencrypt_value(v, key=fld)
304 if new_v != v:
305 new_settings[fld] = new_v
306 report.migrated += 1
307 changed = True
308 else:
309 # Defensive: a legacy value that did not re-encrypt is still
310 # residual, not migrated.
311 report.legacy += 1
312 else:
313 report.plaintext += 1
315 if changed and not dry_run: 315 ↛ 316line 315 didn't jump to line 316 because the condition on line 315 was never true
316 await prisma_client.db.litellm_ssoconfig.update(
317 where={"id": "sso_config"},
318 data={"sso_settings": json.dumps(new_settings)},
319 )
320 return report
323async def _migrate_callback_vars_table(
324 prisma_client: object,
325 table_name: Literal["team", "verification_token"],
326 dry_run: bool,
327) -> LocationReport:
328 """Migrate callback-var credentials on the team or verification-token table.
330 Covers both shapes the ``decrypt_callback_vars`` / ``encrypt_callback_vars``
331 transforms understand: ``metadata.logging[*].callback_vars.<sensitive>`` and
332 the top-level ``metadata.callback_settings.callback_vars.<sensitive>``. Reuses
333 those proven transforms (selective, prefix-marked; legacy plaintext is left
334 alone until re-encrypted).
335 """
336 from litellm.proxy.common_utils.callback_utils import (
337 decrypt_callback_vars,
338 encrypt_callback_vars,
339 )
341 report: Final = LocationReport(location=f"{table_name}.callback_vars")
343 if table_name == "team":
344 table = prisma_client.db.litellm_teamtable
345 pk = "team_id"
346 else:
347 table = prisma_client.db.litellm_verificationtoken
348 pk = "token"
350 rows: Final = await table.find_many()
351 for row in rows or []:
352 metadata = getattr(row, "metadata", None)
353 if not isinstance(metadata, dict) or ("logging" not in metadata and "callback_settings" not in metadata):
354 continue
356 # Classify every callback-var value directly (strip the litellm_enc::
357 # marker, then prefix/decrypt-classify), exactly like the covered-table
358 # scanner. Detecting legacy this way is independent of the AES gate, so
359 # the check_encryption (dry-run) attestation is correct even when run
360 # before the gate is enabled -- a re-encrypt-delta heuristic would read
361 # zero residual here with the gate off.
362 row_legacy = 0
363 for cvs in _iter_callback_var_dicts(metadata):
364 for v in cvs.values(): 364 ↛ 365line 364 didn't jump to line 365 because the loop on line 364 never started
365 report.scanned += 1
366 cls = _classify_callback_value(v)
367 if cls == "migrated":
368 report.already_v2 += 1
369 elif cls == "legacy":
370 row_legacy += 1
371 else: # plaintext / not-a-string
372 report.plaintext += 1
374 if row_legacy == 0: 374 ↛ 377line 374 didn't jump to line 377 because the condition on line 374 was always true
375 continue # no legacy ciphertext in this row
377 if dry_run:
378 # Residual for the attestation; a dry run writes nothing.
379 report.legacy += row_legacy
380 continue
382 # Real run: re-encrypt the legacy ciphertext to AES via the proven
383 # selective transforms and persist. Never drop a row on failure.
384 try:
385 re_encrypted = encrypt_callback_vars(decrypt_callback_vars(metadata))
386 except Exception as e: # pragma: no cover - defensive; never drop a row
387 verbose_proxy_logger.warning(
388 "Skipping %s row %s callback_vars (transform failed): %s",
389 table_name,
390 getattr(row, pk, "?"),
391 str(e),
392 )
393 report.undecryptable += row_legacy
394 continue
395 report.migrated += row_legacy
396 await table.update(
397 where={pk: getattr(row, pk)},
398 data={"metadata": json.dumps(re_encrypted)},
399 )
401 return report
404def _iter_callback_var_dicts(metadata: dict[str, object]):
405 """Yield each ``callback_vars`` dict in a metadata structure.
407 Mirrors ``_transform_callback_vars``: credentials live both under
408 ``logging[*].callback_vars`` and under the top-level
409 ``callback_settings.callback_vars``. Counting only the former would let the
410 walker report success while leaving ``callback_settings`` secrets in legacy
411 format at rest.
412 """
413 for entry in metadata.get("logging", []) or []:
414 if isinstance(entry, dict): 414 ↛ 413line 414 didn't jump to line 413 because the condition on line 414 was always true
415 cvs = entry.get("callback_vars")
416 if isinstance(cvs, dict): 416 ↛ 413line 416 didn't jump to line 413 because the condition on line 416 was always true
417 yield cvs
418 callback_settings: Final = metadata.get("callback_settings")
419 if isinstance(callback_settings, dict):
420 cvs = callback_settings.get("callback_vars")
421 if isinstance(cvs, dict): 421 ↛ exitline 421 didn't return from function '_iter_callback_var_dicts' because the condition on line 421 was always true
422 yield cvs
425def _classify_callback_value(value: object) -> ValueClass:
426 """Classify one stored callback-var value, independent of the AES gate.
428 Encrypted callback vars carry the ``litellm_enc::`` marker in front of the
429 ciphertext; strip it, then classify the inner value the same way the
430 covered-table scanner does (``v2:gcm:`` prefix -> migrated, nacl-decryptable
431 -> legacy, otherwise plaintext). Detecting legacy by decrypt rather than by a
432 re-encrypt delta is what makes the ``check_encryption`` attestation correct
433 even when run with the AES write gate off.
434 """
435 from litellm.proxy.common_utils.callback_utils import (
436 _CALLBACK_VAR_ENCRYPTED_PREFIX,
437 )
439 if not isinstance(value, str):
440 return "not-a-string"
441 inner = value
442 inner = inner.removeprefix(_CALLBACK_VAR_ENCRYPTED_PREFIX)
443 return classify_value(inner, key="callback")
446# ---------------------------------------------------------------------------
447# Read-only scanner for the rotation-covered tables.
448#
449# ``_rotate_master_key`` re-encrypts these tables but returns no counts, so on
450# its own it can neither attest residual legacy nor report how many rows it
451# migrated. This scanner reads (never writes) the same encrypted columns the
452# rotation path touches and classifies every value, giving both the attestation
453# coverage and the pre/post counts the rotation path can't supply itself.
454# ---------------------------------------------------------------------------
456# (location, prisma db attribute, JSON columns to walk, scalar string columns).
457_COVERED_TABLE_SPECS: Final = [
458 ("model_table", "litellm_proxymodeltable", ("litellm_params",), ()),
459 ("credentials", "litellm_credentialstable", ("credential_values",), ()),
460 ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()),
461 ("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)),
462 ("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)),
463]
466def _iter_encrypted_strings(obj: object):
467 """Yield every string leaf in a nested dict/list/scalar structure.
469 Iterative (explicit stack) on purpose: recursion here is banned by the
470 code-quality recursive-function detector (unbounded nesting has caused CPU
471 spikes in the past), and an explicit stack walks arbitrary depth safely.
472 """
473 stack: Final[list[object]] = [obj]
474 while stack:
475 cur = stack.pop()
476 if isinstance(cur, str):
477 yield cur
478 elif isinstance(cur, dict):
479 stack.extend(cur.values())
480 elif isinstance(cur, list):
481 stack.extend(cur)
484def _classify_into_report(report: LocationReport, value: str) -> None:
485 """Classify one stored string and bump the matching read-only counter.
487 Only genuine nacl ciphertext lands in ``legacy``; non-secret strings (model
488 names, base URLs, …) do not decrypt and fall through to ``plaintext``, so
489 over-scanning a column is harmless to the residual count.
490 """
491 report.count(classify_value(value, key="scan"))
494def _classify_secret_map(value: object, key: str) -> ValueClass | None:
495 try:
496 decoded: Final = decode_secret_map(value, key=key)
497 except SecretMapDecodeError:
498 return "undecryptable"
499 if not decoded:
500 return None
501 ciphertext: Final = json.loads(value) if isinstance(value, str) and value.lstrip().startswith('"') else value
502 return "migrated" if is_migrated(ciphertext) else "legacy"
505async def _scan_one_table(
506 prisma_client: object,
507 location: str,
508 db_attr: str,
509 json_columns: tuple,
510 scalar_columns: tuple,
511) -> LocationReport:
512 report: Final = LocationReport(location=location)
513 table: Final = getattr(prisma_client.db, db_attr, None)
514 if table is None: 514 ↛ 515line 514 didn't jump to line 515 because the condition on line 514 was never true
515 return report
516 try:
517 rows: Final = await table.find_many()
518 except Exception as e: # pragma: no cover - table absent / not migrated
519 verbose_proxy_logger.debug("scan: %s unavailable: %s", location, str(e))
520 return report
521 for row in rows or []:
522 for col in json_columns:
523 raw = getattr(row, col, None)
524 if raw is None:
525 continue
526 if db_attr == "litellm_mcpservertable" and col in ("static_headers", "env"):
527 report.count(_classify_secret_map(raw, col))
528 continue
529 if isinstance(raw, str): 529 ↛ 530line 529 didn't jump to line 530 because the condition on line 529 was never true
530 try:
531 raw = json.loads(raw)
532 except (ValueError, TypeError):
533 pass
534 for s in _iter_encrypted_strings(raw):
535 _classify_into_report(report, s)
536 for col in scalar_columns:
537 v = getattr(row, col, None)
538 if isinstance(v, str): 538 ↛ 536line 538 didn't jump to line 536 because the condition on line 538 was always true
539 _classify_into_report(report, v)
540 return report
543async def _scan_config_env_vars(prisma_client: object) -> LocationReport:
544 """Scan the ``environment_variables`` config row (``param_value`` dict)."""
545 report: Final = LocationReport(location="config_environment_variables")
546 try:
547 record: Final = await prisma_client.db.litellm_config.find_unique(where={"param_name": "environment_variables"})
548 except Exception as e: # pragma: no cover - defensive
549 verbose_proxy_logger.debug("scan: config env vars unavailable: %s", str(e))
550 return report
551 if record is None or record.param_value is None:
552 return report
553 value = record.param_value
554 if isinstance(value, str): 554 ↛ 555line 554 didn't jump to line 555 because the condition on line 554 was never true
555 try:
556 value = json.loads(value)
557 except (ValueError, TypeError):
558 value = {}
559 for s in _iter_encrypted_strings(value): 559 ↛ 560line 559 didn't jump to line 560 because the loop on line 559 never started
560 _classify_into_report(report, s)
561 return report
564async def _scan_covered_tables(prisma_client: object) -> list[LocationReport]:
565 """Read-only classification of every rotation-covered table. No writes."""
566 reports: Final[list[LocationReport]] = []
567 for location, db_attr, json_cols, scalar_cols in _COVERED_TABLE_SPECS:
568 reports.append(await _scan_one_table(prisma_client, location, db_attr, json_cols, scalar_cols))
569 reports.append(await _scan_config_env_vars(prisma_client))
570 return reports
573# ---------------------------------------------------------------------------
574# Orchestrator
575# ---------------------------------------------------------------------------
577# vantage_settings / cloudzero_settings sensitive fields (see *_endpoints.py).
578_VANTAGE_SENSITIVE: Final = ["api_key", "integration_token"]
579_CLOUDZERO_SENSITIVE: Final = ["api_key"]
582async def _migrate_covered_tables(prisma_client: object, user_api_key_dict: object) -> list[LocationReport]:
583 """Re-encrypt the tables already covered by ``_rotate_master_key`` (model
584 table, credentials, MCP credential/env tables, config environment_variables)
585 by running that orchestrator in *same-key* mode. With the AES gate on, the
586 re-encrypt writes land in ``v2:`` format.
588 ``_rotate_master_key`` returns no counts, so we bracket it with read-only
589 scans: the pre-scan's legacy total minus the post-scan's gives the number
590 actually migrated per location, and the post-scan supplies the residual /
591 already-v2 / scanned figures. Returns one report per covered location.
592 """
593 from litellm.proxy.management_endpoints.key_management_endpoints import (
594 _rotate_master_key,
595 )
597 pre: Final = {r.location: r for r in await _scan_covered_tables(prisma_client)}
599 current_key: Final = _get_salt_key()
600 if current_key is None:
601 raise RuntimeError(
602 "Cannot migrate covered tables: no salt key / master key is set. Set LITELLM_SALT_KEY before migrating."
603 )
604 await _rotate_master_key(
605 prisma_client=cast("PrismaClient", prisma_client),
606 user_api_key_dict=cast("UserAPIKeyAuth", user_api_key_dict),
607 current_master_key=current_key,
608 new_master_key=current_key, # same key, algorithm-only switch
609 )
611 post: Final = await _scan_covered_tables(prisma_client)
612 for post_report in post:
613 pre_report = pre.get(post_report.location)
614 pre_legacy = pre_report.legacy if pre_report else 0
615 # Everything that was legacy before and is no longer legacy now was
616 # converted this run.
617 post_report.migrated = max(0, pre_legacy - post_report.legacy)
618 return post
621async def migrate_encryption(
622 prisma_client: object,
623 user_api_key_dict: object,
624 dry_run: bool = False,
625) -> MigrationReport:
626 """Run the full at-rest re-encryption migration.
628 Requires ``general_settings.encryption_algorithm == 'aes-256-gcm'`` so writes
629 are produced in the AES format. Idempotent and resumable: re-running skips
630 already-migrated values and finishes any partial run.
632 A ``dry_run`` performs no writes: the covered tables are scanned read-only
633 (so their residual legacy still counts toward the attestation) and the
634 net-new walkers run in dry-run mode.
635 """
636 _assert_aes_gate_enabled()
638 report: Final = MigrationReport()
640 # Tables that already have a rotation path (items 1, 2, 5-10). On a real run
641 # delegate to the rotation path (with bracketing scans for counts); on a dry
642 # run only classify them read-only.
643 if dry_run:
644 for covered in await _scan_covered_tables(prisma_client):
645 report.add(covered)
646 else:
647 for covered in await _migrate_covered_tables(prisma_client, user_api_key_dict):
648 report.add(covered)
650 # Net-new walkers (items 3, 4, 11, 12, 13).
651 report.add(await _migrate_callback_vars_table(prisma_client, "team", dry_run))
652 report.add(await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run))
653 report.add(await _migrate_config_settings_row(prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run))
654 report.add(await _migrate_config_settings_row(prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run))
655 report.add(await _migrate_sso_config(prisma_client, dry_run))
657 return report
660async def check_encryption(prisma_client: object) -> MigrationReport:
661 """Read-only residual scan across **every** at-rest location. No writes.
663 Covers both the rotation-managed tables (model / credentials / MCP credential
664 and env-var tables / config ``environment_variables``) and the net-new walker
665 locations (team and verification-token ``callback_vars``, vantage / cloudzero
666 config rows, SSO config). Reports how many values are still ``legacy``;
667 ``residual_legacy == 0`` across this full scan is the compliance attestation.
668 """
669 report: Final = MigrationReport()
671 # Rotation-covered tables (read-only classification).
672 for covered in await _scan_covered_tables(prisma_client):
673 report.add(covered)
675 # Net-new walker locations, in dry-run (read-only) mode.
676 report.add(await _migrate_callback_vars_table(prisma_client, "team", dry_run=True))
677 report.add(await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run=True))
678 report.add(await _migrate_config_settings_row(prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run=True))
679 report.add(
680 await _migrate_config_settings_row(prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run=True)
681 )
682 report.add(await _migrate_sso_config(prisma_client, dry_run=True))
683 return report