mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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:
parent
0eb7d975a1
commit
50606bf090
2 changed files with 39 additions and 4 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue