ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

UI SSO - fix reset env var when ui_access_mode is updated (#13011)

Browse Source 
* fix(ui_sso.py): fix form action on login when sso is enabled

* fix: multiple fixes - fix resetting env var in proxy config + add key to exception message on key decryption

fixes issue where env vars would be reset

* refactor(proxy_server.py): cleanup redundant decryption line

* fix(proxy_setting_endpoints.py): show saved ui access mode

allows admin to know what they'd previously stored in db
This commit is contained in:
Krish Dholakia 2025-07-26 11:42:41 -07:00 committed by GitHub
parent 4b50566d6d
commit eed0cf2ee9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 206 additions and 166 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
litellm/proxy/auth/auth_checks.py
View File

@ -1025,7 +1025,9 @@ class ExperimentalUIJWTToken:
decrypt_value_helper,
)
decrypted_token = decrypt_value_helper(hashed_token, exception_type="debug")
decrypted_token = decrypt_value_helper(
hashed_token, key="ui_hash_key", exception_type="debug"
)
if decrypted_token is None:
return None
try:

10
litellm/proxy/common_utils/encrypt_decrypt_utils.py
View File

@ -40,7 +40,9 @@ def encrypt_value_helper(value: str, new_encryption_key: Optional[str] = None):
def decrypt_value_helper(
value: str, exception_type: Literal["debug", "error"] = "error"
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",
):
signing_key = _get_salt_key()
@ -53,11 +55,15 @@ def decrypt_value_helper(
# if it's not str - do not decrypt it, return the value
return value
except Exception as e:
error_message = f"Error decrypting value, 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"
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 None
verbose_proxy_logger.debug(
f"Unable to decrypt value={value} for key: {key}, returning None"
)
verbose_proxy_logger.exception(error_message)
# [Non-Blocking Exception. - this should not block decrypting other values]
return None

41
litellm/proxy/management_endpoints/ui_sso.py
View File

@ -58,7 +58,12 @@ from litellm.proxy.management_endpoints.sso_helper_utils import (
)
from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add
from litellm.proxy.management_endpoints.types import CustomOpenID
from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path
from litellm.proxy.utils import (
PrismaClient,
ProxyLogging,
get_custom_url,
get_server_root_path,
)
from litellm.secret_managers.main import get_secret_bool, str_to_bool
from litellm.types.proxy.management_endpoints.ui_sso import *
@ -162,14 +167,8 @@ async def serve_login_page(
</div>
"""
# Get the base URL for form action - CHANGE THIS TO POINT TO /login
proxy_base_url = os.getenv("PROXY_BASE_URL", "")
server_root_path = os.getenv("SERVER_ROOT_PATH", "")
if server_root_path != "":
proxy_base_url += server_root_path
form_action = (
proxy_base_url + "/sso/key/generate"
) # CHANGE BACK to /sso/key/generate
# Get the base URL for form action using proper URL construction
form_action = get_custom_url(request_base_url=str(request.base_url), route="login")
unified_login_html = f"""
<!DOCTYPE html>
@ -1045,9 +1044,9 @@ async def insert_sso_user(
if user_defined_values.get("max_budget") is None:
user_defined_values["max_budget"] = litellm.max_internal_user_budget
if user_defined_values.get("budget_duration") is None:
user_defined_values[
"budget_duration"
] = litellm.internal_user_budget_duration
user_defined_values["budget_duration"] = (
litellm.internal_user_budget_duration
)
if user_defined_values["user_role"] is None:
user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
@ -1246,9 +1245,9 @@ class SSOAuthenticationHandler:
if state:
redirect_params["state"] = state
elif "okta" in generic_authorization_endpoint:
redirect_params[
"state"
] = uuid.uuid4().hex # set state param for okta - required
redirect_params["state"] = (
uuid.uuid4().hex
) # set state param for okta - required
return await generic_sso.get_login_redirect(**redirect_params) # type: ignore
raise ValueError(
"Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso"
@ -1769,9 +1768,9 @@ class MicrosoftSSOHandler:
# if user is trying to get the raw sso response for debugging, return the raw sso response
if return_raw_sso_response:
original_msft_result[
MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY
] = user_team_ids
original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = (
user_team_ids
)
return original_msft_result or {}
result = MicrosoftSSOHandler.openid_from_response(
@ -1839,9 +1838,9 @@ class MicrosoftSSOHandler:
# Fetch user membership from Microsoft Graph API
all_group_ids = []
next_link: Optional[
str
] = MicrosoftSSOHandler.graph_api_user_groups_endpoint
next_link: Optional[str] = (
MicrosoftSSOHandler.graph_api_user_groups_endpoint
)
auth_headers = {"Authorization": f"Bearer {access_token}"}
page_count = 0

37
litellm/proxy/proxy_server.py
View File

@ -2422,7 +2422,7 @@ class ProxyConfig:
for k, v in _litellm_params.items():
if isinstance(v, str):
# decrypt value
_value = decrypt_value_helper(value=v)
_value = decrypt_value_helper(value=v, key=k)
if _value is None:
raise Exception("Unable to decrypt value={}".format(v))
# sanity check if string > size 0
@ -2458,7 +2458,7 @@ class ProxyConfig:
if isinstance(_litellm_params, dict):
# decrypt values
for k, v in _litellm_params.items():
decrypted_value = decrypt_value_helper(value=v)
decrypted_value = decrypt_value_helper(value=v, key=k)
_litellm_params[k] = decrypted_value
_litellm_params = LiteLLM_Params(**_litellm_params)
else:
@ -2522,9 +2522,6 @@ class ProxyConfig:
config_data = await proxy_config.get_config()
self._add_callbacks_from_db_config(config_data)
# we need to set env variables too
self._add_environment_variables_from_db_config(config_data)
# router settings
await self._add_router_settings_from_db_config(
config_data=config_data, llm_router=llm_router, prisma_client=prisma_client
@ -2574,13 +2571,6 @@ class ProxyConfig:
failure_callback
)
def _add_environment_variables_from_db_config(self, config_data: dict) -> None:
"""
Adds environment variables from DB config to litellm
"""
environment_variables = config_data.get("environment_variables", {})
self._decrypt_and_set_db_env_variables(environment_variables)
def _encrypt_env_variables(
self, environment_variables: dict, new_encryption_key: Optional[str] = None
) -> dict:
@ -2606,7 +2596,7 @@ class ProxyConfig:
decrypted_env_vars = {}
for k, v in environment_variables.items():
try:
decrypted_value = decrypt_value_helper(value=v)
decrypted_value = decrypt_value_helper(value=v, key=k)
if decrypted_value is not None:
os.environ[k] = decrypted_value
decrypted_env_vars[k] = decrypted_value
@ -2759,7 +2749,10 @@ class ProxyConfig:
"""
if param_name == "environment_variables":
self._decrypt_and_set_db_env_variables(db_param_value)
decrypted_env_vars = self._decrypt_and_set_db_env_variables(db_param_value)
current_config.setdefault("environment_variables", {}).update(
decrypted_env_vars
)
return current_config
elif param_name == "litellm_settings" and isinstance(db_param_value, dict):
for key, value in db_param_value.items():
@ -2991,7 +2984,7 @@ class ProxyConfig:
decrypted_credential_values = {}
for k, v in credential_object.credential_values.items():
decrypted_credential_values[k] = decrypt_value_helper(v) or v
decrypted_credential_values[k] = decrypt_value_helper(value=v, key=k) or v
credential_object.credential_values = decrypted_credential_values
return credential_object
@ -8614,7 +8607,9 @@ async def get_config(): # noqa: PLR0915
env_vars_dict[_var] = None
else:
# decode + decrypt the value
decrypted_value = decrypt_value_helper(value=env_variable)
decrypted_value = decrypt_value_helper(
value=env_variable, key=_var
)
env_vars_dict[_var] = decrypted_value
_data_to_return.append({"name": _callback, "variables": env_vars_dict})
@ -8631,7 +8626,9 @@ async def get_config(): # noqa: PLR0915
_langfuse_env_vars[_var] = None
else:
# decode + decrypt the value
decrypted_value = decrypt_value_helper(value=env_variable)
decrypted_value = decrypt_value_helper(
value=env_variable, key=_var
)
_langfuse_env_vars[_var] = decrypted_value
_data_to_return.append(
@ -8653,7 +8650,9 @@ async def get_config(): # noqa: PLR0915
_slack_env_vars[_var] = _value
else:
# decode + decrypt the value
_decrypted_value = decrypt_value_helper(value=env_variable)
_decrypted_value = decrypt_value_helper(
value=env_variable, key=_var
)
_slack_env_vars[_var] = _decrypted_value
_alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types
@ -8689,7 +8688,7 @@ async def get_config(): # noqa: PLR0915
_email_env_vars[_var] = None
else:
# decode + decrypt the value
_decrypted_value = decrypt_value_helper(value=env_variable)
_decrypted_value = decrypt_value_helper(value=env_variable, key=_var)
_email_env_vars[_var] = _decrypted_value
alerting_data.append(

240
litellm/proxy/spend_tracking/cloudzero_endpoints.py
View File

@ -29,29 +29,29 @@ _sensitive_masker = SensitiveDataMasker()
async def _set_cloudzero_settings(api_key: str, connection_id: str, timezone: str):
"""
Store CloudZero settings in the database with encrypted API key.
Args:
api_key: CloudZero API key to encrypt and store
connection_id: CloudZero connection ID
timezone: Timezone for date handling
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
# Encrypt the API key before storing
encrypted_api_key = encrypt_value_helper(api_key)
cloudzero_settings = {
"api_key": encrypted_api_key,
"connection_id": connection_id,
"timezone": timezone,
}
await prisma_client.db.litellm_config.upsert(
where={"param_name": "cloudzero_settings"},
data={
@ -67,41 +67,47 @@ async def _set_cloudzero_settings(api_key: str, connection_id: str, timezone: st
async def _get_cloudzero_settings():
"""
Retrieve CloudZero settings from the database with decrypted API key.
Returns:
dict: CloudZero settings with decrypted API key
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
cloudzero_config = await prisma_client.db.litellm_config.find_first(
where={"param_name": "cloudzero_settings"}
)
if not cloudzero_config or not cloudzero_config.param_value:
raise HTTPException(
status_code=400,
detail={"error": "CloudZero settings not configured. Please run /cloudzero/init first."}
detail={
"error": "CloudZero settings not configured. Please run /cloudzero/init first."
},
)
settings = dict(cloudzero_config.param_value)
# Decrypt the API key
encrypted_api_key = settings.get("api_key")
if encrypted_api_key:
decrypted_api_key = decrypt_value_helper(encrypted_api_key, exception_type="error")
decrypted_api_key = decrypt_value_helper(
encrypted_api_key, key="cloudzero_api_key", exception_type="error"
)
if decrypted_api_key is None:
raise HTTPException(
status_code=500,
detail={"error": "Failed to decrypt CloudZero API key. Check your salt key configuration."}
detail={
"error": "Failed to decrypt CloudZero API key. Check your salt key configuration."
},
)
settings["api_key"] = decrypted_api_key
return settings
@ -116,10 +122,10 @@ async def get_cloudzero_settings(
):
"""
View current CloudZero settings.
Returns the current CloudZero configuration with the API key masked for security.
Only the first 4 and last 4 characters of the API key are shown.
Only admin users can view CloudZero settings.
"""
# Validation
@ -128,34 +134,33 @@ async def get_cloudzero_settings(
status_code=403,
detail={"error": CommonProxyErrors.not_allowed_access.value},
)
try:
# Get CloudZero settings using the accessor method
settings = await _get_cloudzero_settings()
# Use SensitiveDataMasker to mask the API key
masked_settings = _sensitive_masker.mask_dict(settings)
return CloudZeroSettingsView(
api_key_masked=masked_settings["api_key"],
connection_id=settings["connection_id"],
timezone=settings["timezone"],
status="configured"
status="configured",
)
except HTTPException as e:
if e.status_code == 400:
# Settings not configured
raise HTTPException(
status_code=404,
detail={"error": "CloudZero settings not configured"}
status_code=404, detail={"error": "CloudZero settings not configured"}
)
raise e
except Exception as e:
verbose_proxy_logger.error(f"Error retrieving CloudZero settings: {str(e)}")
raise HTTPException(
status_code=500,
detail={"error": f"Failed to retrieve CloudZero settings: {str(e)}"}
detail={"error": f"Failed to retrieve CloudZero settings: {str(e)}"},
)
@ -171,15 +176,15 @@ async def update_cloudzero_settings(
):
"""
Update existing CloudZero settings.
Allows updating individual CloudZero configuration fields without requiring all fields.
Only provided fields will be updated; others will remain unchanged.
Parameters:
- api_key: (Optional) New CloudZero API key for authentication
- connection_id: (Optional) New CloudZero connection ID for data submission
- timezone: (Optional) New timezone for date handling
Only admin users can update CloudZero settings.
"""
# Validation
@ -188,52 +193,66 @@ async def update_cloudzero_settings(
status_code=403,
detail={"error": CommonProxyErrors.not_allowed_access.value},
)
# Check if at least one field is provided
if not any([request.api_key, request.connection_id, request.timezone]):
raise HTTPException(
status_code=400,
detail={"error": "At least one field must be provided for update"}
detail={"error": "At least one field must be provided for update"},
)
try:
# Get current settings
current_settings = await _get_cloudzero_settings()
# Update only provided fields
updated_api_key = request.api_key if request.api_key is not None else current_settings["api_key"]
updated_connection_id = request.connection_id if request.connection_id is not None else current_settings["connection_id"]
updated_timezone = request.timezone if request.timezone is not None else current_settings["timezone"]
updated_api_key = (
request.api_key
if request.api_key is not None
else current_settings["api_key"]
)
updated_connection_id = (
request.connection_id
if request.connection_id is not None
else current_settings["connection_id"]
)
updated_timezone = (
request.timezone
if request.timezone is not None
else current_settings["timezone"]
)
# Store updated settings using the setter method with encryption
await _set_cloudzero_settings(
api_key=updated_api_key,
connection_id=updated_connection_id,
timezone=updated_timezone
timezone=updated_timezone,
)
verbose_proxy_logger.info("CloudZero settings updated successfully")
return CloudZeroInitResponse(
message="CloudZero settings updated successfully",
status="success"
message="CloudZero settings updated successfully", status="success"
)
except HTTPException as e:
if e.status_code == 400:
# Settings not configured yet
raise HTTPException(
status_code=404,
detail={"error": "CloudZero settings not found. Please initialize settings first using /cloudzero/init"}
detail={
"error": "CloudZero settings not found. Please initialize settings first using /cloudzero/init"
},
)
raise e
except Exception as e:
verbose_proxy_logger.error(f"Error updating CloudZero settings: {str(e)}")
raise HTTPException(
status_code=500,
detail={"error": f"Failed to update CloudZero settings: {str(e)}"}
detail={"error": f"Failed to update CloudZero settings: {str(e)}"},
)
# Global variable to track if CloudZero background job has been initialized
_cloudzero_background_job_initialized = False
@ -244,77 +263,86 @@ async def init_cloudzero_background_job():
This should be called from the proxy server startup.
"""
global _cloudzero_background_job_initialized
if _cloudzero_background_job_initialized:
verbose_proxy_logger.debug("CloudZero background job already initialized, skipping")
verbose_proxy_logger.debug(
"CloudZero background job already initialized, skipping"
)
return
try:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
verbose_proxy_logger.warning("Prisma client not available, skipping CloudZero background job initialization")
verbose_proxy_logger.warning(
"Prisma client not available, skipping CloudZero background job initialization"
)
return
# Get CloudZero settings from database
cloudzero_config = await prisma_client.db.litellm_config.find_first(
where={"param_name": "cloudzero_settings"}
)
if not cloudzero_config or not cloudzero_config.param_value:
verbose_proxy_logger.debug("CloudZero settings not configured, skipping background job initialization")
verbose_proxy_logger.debug(
"CloudZero settings not configured, skipping background job initialization"
)
return
settings = dict(cloudzero_config.param_value)
# Initialize CloudZero logger with credentials
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
logger = CloudZeroLogger(
api_key=settings["api_key"],
connection_id=settings["connection_id"],
timezone=settings["timezone"]
timezone=settings["timezone"],
)
# Initialize the background job
await logger.init_background_job()
_cloudzero_background_job_initialized = True
verbose_proxy_logger.info("CloudZero background job initialized successfully")
except Exception as e:
verbose_proxy_logger.error(f"Error initializing CloudZero background job: {str(e)}")
verbose_proxy_logger.error(
f"Error initializing CloudZero background job: {str(e)}"
)
async def is_cloudzero_setup_in_db() -> bool:
"""
Check if CloudZero is setup in the database.
CloudZero is considered setup in the database if:
- CloudZero settings exist in the database
- CloudZero settings exist in the database
- The settings have a non-None value
Returns:
bool: True if CloudZero is active, False otherwise
"""
try:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return False
# Check for CloudZero settings in database
cloudzero_config = await prisma_client.db.litellm_config.find_first(
where={"param_name": "cloudzero_settings"}
)
# CloudZero is setup in the database if config exists and has non-None value
return cloudzero_config is not None and cloudzero_config.param_value is not None
except Exception as e:
verbose_proxy_logger.error(f"Error checking CloudZero status: {str(e)}")
return False
@router.post(
"/cloudzero/init",
tags=["CloudZero"],
@ -327,15 +355,15 @@ async def init_cloudzero_settings(
):
"""
Initialize CloudZero settings and store in the database.
This endpoint stores the CloudZero API key, connection ID, and timezone configuration
in the proxy database for use by the CloudZero logger.
Parameters:
- api_key: CloudZero API key for authentication
- connection_id: CloudZero connection ID for data submission
- connection_id: CloudZero connection ID for data submission
- timezone: Timezone for date handling (default: UTC)
Only admin users can configure CloudZero settings.
"""
# Validation
@ -344,30 +372,29 @@ async def init_cloudzero_settings(
status_code=403,
detail={"error": CommonProxyErrors.not_allowed_access.value},
)
try:
# Store settings using the setter method with encryption
await _set_cloudzero_settings(
api_key=request.api_key,
connection_id=request.connection_id,
timezone=request.timezone
timezone=request.timezone,
)
verbose_proxy_logger.info("CloudZero settings initialized successfully")
# Initialize background job after settings are saved
await init_cloudzero_background_job()
return CloudZeroInitResponse(
message="CloudZero settings initialized successfully",
status="success"
message="CloudZero settings initialized successfully", status="success"
)
except Exception as e:
verbose_proxy_logger.error(f"Error initializing CloudZero settings: {str(e)}")
raise HTTPException(
status_code=500,
detail={"error": f"Failed to initialize CloudZero settings: {str(e)}"}
detail={"error": f"Failed to initialize CloudZero settings: {str(e)}"},
)
@ -383,13 +410,13 @@ async def cloudzero_dry_run_export(
):
"""
Perform a dry run export using the CloudZero logger.
This endpoint uses the CloudZero logger to perform a dry run export,
which displays the data that would be exported without actually sending it to CloudZero.
Parameters:
- limit: Optional limit on number of records to process (default: 10000)
Only admin users can perform CloudZero exports.
"""
from datetime import datetime
@ -400,27 +427,31 @@ async def cloudzero_dry_run_export(
status_code=403,
detail={"error": CommonProxyErrors.not_allowed_access.value},
)
try:
# Import and initialize CloudZero logger with credentials
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
# Initialize logger with credentials directly
logger = CloudZeroLogger()
await logger.dry_run_export_usage_data(target_hour=datetime.utcnow(), limit=request.limit)
await logger.dry_run_export_usage_data(
target_hour=datetime.utcnow(), limit=request.limit
)
verbose_proxy_logger.info("CloudZero dry run export completed successfully")
return CloudZeroExportResponse(
message="CloudZero dry run export completed successfully. Check logs for output.",
status="success"
status="success",
)
except Exception as e:
verbose_proxy_logger.error(f"Error performing CloudZero dry run export: {str(e)}")
verbose_proxy_logger.error(
f"Error performing CloudZero dry run export: {str(e)}"
)
raise HTTPException(
status_code=500,
detail={"error": f"Failed to perform CloudZero dry run export: {str(e)}"}
detail={"error": f"Failed to perform CloudZero dry run export: {str(e)}"},
)
@ -436,13 +467,13 @@ async def cloudzero_export(
):
"""
Perform an actual export using the CloudZero logger.
This endpoint uses the CloudZero logger to export usage data to CloudZero AnyCost API.
Parameters:
- limit: Optional limit on number of records to export
- operation: CloudZero operation type ("replace_hourly" or "sum", default: "replace_hourly")
Only admin users can perform CloudZero exports.
"""
@ -453,11 +484,11 @@ async def cloudzero_export(
status_code=403,
detail={"error": CommonProxyErrors.not_allowed_access.value},
)
try:
# Get CloudZero settings using the accessor method with decryption
settings = await _get_cloudzero_settings()
# Import and initialize CloudZero logger with credentials
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
@ -465,26 +496,23 @@ async def cloudzero_export(
logger = CloudZeroLogger(
api_key=settings["api_key"],
connection_id=settings["connection_id"],
timezone=settings["timezone"]
timezone=settings["timezone"],
)
await logger.export_usage_data(
target_hour=datetime.utcnow(),
limit=request.limit,
operation=request.operation
target_hour=datetime.utcnow(),
limit=request.limit,
operation=request.operation,
)
verbose_proxy_logger.info("CloudZero export completed successfully")
return CloudZeroExportResponse(
message="CloudZero export completed successfully",
status="success"
message="CloudZero export completed successfully", status="success"
)
except Exception as e:
verbose_proxy_logger.error(f"Error performing CloudZero export: {str(e)}")
raise HTTPException(
status_code=500,
detail={"error": f"Failed to perform CloudZero export: {str(e)}"}
detail={"error": f"Failed to perform CloudZero export: {str(e)}"},
)

8
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View File

@ -401,6 +401,7 @@ async def get_sso_settings():
generic_userinfo_endpoint=get_env_value("GENERIC_USERINFO_ENDPOINT"),
proxy_base_url=get_env_value("PROXY_BASE_URL"),
user_email=proxy_admin_email, # Get from config instead of environment
ui_access_mode=general_settings.get("ui_access_mode", None),
)
# Get the schema for UI display
@ -471,7 +472,6 @@ async def update_sso_settings(sso_config: SSOConfig):
# Update environment variables in config and in memory
sso_data = sso_config.model_dump(exclude_none=True)
mapped_env_vars = {}
for field_name, value in sso_data.items():
if field_name == "user_email" and value is not None:
@ -484,14 +484,14 @@ async def update_sso_settings(sso_config: SSOConfig):
env_var_name = env_var_mapping[field_name]
# Update in config
config["environment_variables"][env_var_name] = value
mapped_env_vars[env_var_name] = value
# Update in runtime environment
os.environ[env_var_name] = value
stored_config = config
if len(mapped_env_vars) > 0:
if len(config["environment_variables"]) > 0:
stored_config["environment_variables"] = proxy_config._encrypt_env_variables(
environment_variables=mapped_env_vars
environment_variables=config["environment_variables"]
)
# Save the updated config
await proxy_config.save_config(new_config=stored_config)

28
tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py
View File

@ -21,26 +21,30 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
def test_encrypt_decrypt_with_master_key():
setattr(proxy_server, "master_key", "sk-1234")
assert decrypt_value_helper(encrypt_value_helper("test")) == "test"
assert decrypt_value_helper(encrypt_value_helper(10)) == 10
assert decrypt_value_helper(encrypt_value_helper(True)) is True
assert decrypt_value_helper(encrypt_value_helper(None)) is None
assert decrypt_value_helper(encrypt_value_helper({"rpm": 10})) == {"rpm": 10}
assert decrypt_value_helper(encrypt_value_helper("test", key="test_key")) == "test"
assert decrypt_value_helper(encrypt_value_helper(10, key="test_key")) == 10
assert decrypt_value_helper(encrypt_value_helper(True, key="test_key")) is True
assert decrypt_value_helper(encrypt_value_helper(None, key="test_key")) is None
assert decrypt_value_helper(encrypt_value_helper({"rpm": 10}, key="test_key")) == {
"rpm": 10
}
# encryption should actually occur for strings
assert encrypt_value_helper("test") != "test"
assert encrypt_value_helper("test", key="test_key") != "test"
def test_encrypt_decrypt_with_salt_key():
os.environ["LITELLM_SALT_KEY"] = "sk-salt-key2222"
print(f"LITELLM_SALT_KEY: {os.environ['LITELLM_SALT_KEY']}")
assert decrypt_value_helper(encrypt_value_helper("test")) == "test"
assert decrypt_value_helper(encrypt_value_helper(10)) == 10
assert decrypt_value_helper(encrypt_value_helper(True)) is True
assert decrypt_value_helper(encrypt_value_helper(None)) is None
assert decrypt_value_helper(encrypt_value_helper({"rpm": 10})) == {"rpm": 10}
assert decrypt_value_helper(encrypt_value_helper("test", key="test_key")) == "test"
assert decrypt_value_helper(encrypt_value_helper(10, key="test_key")) == 10
assert decrypt_value_helper(encrypt_value_helper(True, key="test_key")) is True
assert decrypt_value_helper(encrypt_value_helper(None, key="test_key")) is None
assert decrypt_value_helper(encrypt_value_helper({"rpm": 10}, key="test_key")) == {
"rpm": 10
}
# encryption should actually occur for strings
assert encrypt_value_helper("test") != "test"
assert encrypt_value_helper("test", key="test_key") != "test"
os.environ.pop("LITELLM_SALT_KEY", None)

4
tests/test_litellm/proxy/auth/test_auth_checks.py
View File

@ -68,7 +68,9 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v
)
# Decrypt and verify token contents
decrypted_token = decrypt_value_helper(token, exception_type="debug")
decrypted_token = decrypt_value_helper(
token, key="ui_hash_key", exception_type="debug"
)
# Check that decrypted_token is not None before using json.loads
assert decrypted_token is not None
token_data = json.loads(decrypted_token)