Added new step into rotate master key function for processing credentials table (#17952)

* fix: Return 403 exception when calling GET responses api

* fix: added new step into rotate master key function for processing credentials table
This commit is contained in:
Eric84626 2025-12-16 12:10:22 +08:00 committed by GitHub
parent 0eb7d975a1
commit 50606bf090
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 39 additions and 4 deletions

View file

@ -21,11 +21,11 @@ router = APIRouter()
class CredentialHelperUtils:
@staticmethod
def encrypt_credential_values(credential: CredentialItem) -> CredentialItem:
def encrypt_credential_values(credential: CredentialItem, new_encryption_key: Optional[str] = None) -> CredentialItem:
"""Encrypt values in credential.credential_values and add to DB"""
encrypted_credential_values = {}
for key, value in (credential.credential_values or {}).items():
encrypted_credential_values[key] = encrypt_value_helper(value)
encrypted_credential_values[key] = encrypt_value_helper(value, new_encryption_key)
# Return a new object to avoid mutating the caller's credential, which
# is kept in memory and should remain unencrypted.
@ -246,7 +246,7 @@ async def delete_credential(
def update_db_credential(
db_credential: CredentialItem, updated_patch: CredentialItem
db_credential: CredentialItem, updated_patch: CredentialItem, new_encryption_key: Optional[str] = None
) -> CredentialItem:
"""
Update a credential in the DB.
@ -258,7 +258,8 @@ def update_db_credential(
)
encrypted_credential = CredentialHelperUtils.encrypt_credential_values(
updated_patch
updated_patch,
new_encryption_key,
)
# update model name
if encrypted_credential.credential_name:

View file

@ -2539,6 +2539,40 @@ async def _rotate_master_key(
new_master_key=new_master_key,
)
# 5. process credentials table
try:
credentials = await prisma_client.db.litellm_credentialstable.find_many()
except Exception:
credentials = None
if credentials:
from litellm.proxy.credential_endpoints.endpoints import update_db_credential
for cred in credentials:
try:
decrypted_cred = proxy_config.decrypt_credentials(cred)
encrypted_cred = update_db_credential(
db_credential=cred,
updated_patch=decrypted_cred,
new_encryption_key=new_master_key,
)
credential_object_jsonified = jsonify_object(encrypted_cred.model_dump())
await prisma_client.db.litellm_credentialstable.update(
where={"credential_name": cred.credential_name},
data={
**credential_object_jsonified,
"updated_by": user_api_key_dict.user_id,
},
)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to re-encrypt credential {cred.credential_name}: {str(e)}"
)
# Continue with next credential instead of failing entire rotation
continue
verbose_proxy_logger.debug(
f"Successfully re-encrypted {len(credentials)} credentials with new master key"
)
def get_new_token(data: Optional[RegenerateKeyRequest]) -> str:
if data and data.new_key is not None: