mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
72bbdfd3f3
* feat(responses_id_security.py): encrypt response.id - prevent user A from retrieving user B's response additional security for retrievals on shared accounts Closes LIT-1307 * feat(responses_id_security.py): allow admin to disable responses id security check * test: add initial unit testing * feat(responses_id_security.py): add streaming support * docs: document new param * docs: document new param * feat(responses_id_security.py): add team id checks - ensure it works for service accounts prevent service accounts keys from different teams from accessing each other's responses more secure * test: add unit testing * fix: fix linting error
123 lines
3.8 KiB
Python
123 lines
3.8 KiB
Python
import base64
|
|
import os
|
|
from typing import Literal, Optional
|
|
|
|
from litellm._logging import verbose_proxy_logger
|
|
|
|
|
|
def _get_salt_key():
|
|
from litellm.proxy.proxy_server import master_key
|
|
|
|
salt_key = os.getenv("LITELLM_SALT_KEY", None)
|
|
|
|
if salt_key is None:
|
|
salt_key = master_key
|
|
|
|
return salt_key
|
|
|
|
|
|
def encrypt_value_helper(value: str, new_encryption_key: Optional[str] = None):
|
|
signing_key = new_encryption_key or _get_salt_key()
|
|
|
|
try:
|
|
if isinstance(value, str):
|
|
encrypted_value = encrypt_value(value=value, signing_key=signing_key) # type: ignore
|
|
# Use urlsafe_b64encode for URL-safe base64 encoding (replaces + with - and / with _)
|
|
encrypted_value = base64.urlsafe_b64encode(encrypted_value).decode("utf-8")
|
|
|
|
return encrypted_value
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"Invalid value type passed to encrypt_value: {type(value)} for Value: {value}\n Value must be a string"
|
|
)
|
|
# if it's not a string - do not encrypt it and return the value
|
|
return value
|
|
except Exception as e:
|
|
raise e
|
|
|
|
|
|
def decrypt_value_helper(
|
|
value: str,
|
|
key: str, # this is just for debug purposes, showing the k,v pair that's invalid. not a signing key.
|
|
exception_type: Literal["debug", "error"] = "error",
|
|
return_original_value: bool = False,
|
|
):
|
|
signing_key = _get_salt_key()
|
|
|
|
try:
|
|
if isinstance(value, str):
|
|
# Try URL-safe base64 decoding first (new format)
|
|
# Fall back to standard base64 decoding for backwards compatibility (old format)
|
|
try:
|
|
decoded_b64 = base64.urlsafe_b64decode(value)
|
|
except Exception:
|
|
# If URL-safe decoding fails, try standard base64 decoding for backwards compatibility
|
|
decoded_b64 = base64.b64decode(value)
|
|
|
|
value = decrypt_value(value=decoded_b64, signing_key=signing_key) # type: ignore
|
|
return value
|
|
|
|
# if it's not str - do not decrypt it, return the value
|
|
return value
|
|
except Exception as e:
|
|
|
|
error_message = f"Error decrypting value for key: {key}, Did your master_key/salt key change recently? \nError: {str(e)}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key"
|
|
if exception_type == "debug":
|
|
verbose_proxy_logger.debug(error_message)
|
|
return value if return_original_value else None
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"Unable to decrypt value={value} for key: {key}, returning None"
|
|
)
|
|
if return_original_value:
|
|
return value
|
|
else:
|
|
verbose_proxy_logger.exception(error_message)
|
|
# [Non-Blocking Exception. - this should not block decrypting other values]
|
|
return None
|
|
|
|
|
|
def encrypt_value(value: str, signing_key: str):
|
|
import hashlib
|
|
|
|
import nacl.secret
|
|
import nacl.utils
|
|
|
|
# get 32 byte master key #
|
|
hash_object = hashlib.sha256(signing_key.encode())
|
|
hash_bytes = hash_object.digest()
|
|
|
|
# initialize secret box #
|
|
box = nacl.secret.SecretBox(hash_bytes)
|
|
|
|
# encode message #
|
|
value_bytes = value.encode("utf-8")
|
|
|
|
encrypted = box.encrypt(value_bytes)
|
|
|
|
return encrypted
|
|
|
|
|
|
def decrypt_value(value: bytes, signing_key: str) -> str:
|
|
import hashlib
|
|
|
|
import nacl.secret
|
|
import nacl.utils
|
|
|
|
# get 32 byte master key #
|
|
hash_object = hashlib.sha256(signing_key.encode())
|
|
hash_bytes = hash_object.digest()
|
|
|
|
# initialize secret box #
|
|
box = nacl.secret.SecretBox(hash_bytes)
|
|
|
|
# Convert the bytes object to a string
|
|
try:
|
|
if len(value) == 0:
|
|
return ""
|
|
|
|
plaintext = box.decrypt(value)
|
|
plaintext = plaintext.decode("utf-8") # type: ignore
|
|
return plaintext # type: ignore
|
|
except Exception as e:
|
|
raise e
|