diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5b50ce0730..1e57e5326a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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: diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index c44daafaed..4b71df0117 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -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 diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 25045c1638..c48e1c6cc5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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( """ - # 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""" @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index de410985ac..8dfbf094ef 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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( diff --git a/litellm/proxy/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index c9189cde1d..67de202aa7 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -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)}"}, ) - - diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index a3aafda223..13d35c47ad 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -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) diff --git a/tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py b/tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py index f9c3ff42d0..13c5cc269c 100644 --- a/tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py +++ b/tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py @@ -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) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 635e39c37f..c61d86857f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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)