From d34182833c86a32e8c5a33adaccc489045fd1d5d Mon Sep 17 00:00:00 2001 From: yangdx Date: Wed, 1 Apr 2026 17:32:01 +0800 Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor(health=5Fendpoint?= =?UTF-8?q?s):=20extract=20encryption=20utilities=20and=20update=20secret?= =?UTF-8?q?=20resolution?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - remove inline decrypt_value_helper and _decrypt_value functions from health endpoints - import decrypt_value_helper from common_utils.encrypt_decrypt_utils module - rename get_secret to _get_env_secret to avoid naming conflict with secret_managers.main.get_secret - update _resolve_os_environ_variables to use centralized get_secret from secret_managers.main - add comprehensive tests for _resolve_os_environ_variables with nested dicts and lists --- .../health_endpoints/_health_endpoints.py | 73 +------------------ .../health_endpoints/test_health_endpoints.py | 65 +++++++++++++++++ 2 files changed, 69 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index e07c0ff33b3..86bbfda1a96 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1,11 +1,8 @@ import asyncio -import base64 import copy -import hashlib import json import logging import os -import sys import time import traceback from datetime import datetime, timedelta @@ -30,6 +27,7 @@ from litellm.proxy._types import ( WebhookEvent, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.health_check import ( _clean_endpoint_data, @@ -40,6 +38,7 @@ from litellm.proxy.health_check import ( from litellm.proxy.middleware.in_flight_requests_middleware import ( get_in_flight_requests, ) +from litellm.secret_managers.main import get_secret #### Health ENDPOINTS #### @@ -56,7 +55,7 @@ def _str_to_bool(value: Optional[str]) -> Optional[bool]: return None -def get_secret( +def _get_env_secret( secret_name: str, default_value: Optional[Union[str, bool]] = None ) -> Optional[Union[str, bool]]: if secret_name.startswith("os.environ/"): @@ -72,7 +71,7 @@ def get_secret( def get_secret_bool( secret_name: str, default_value: Optional[bool] = None ) -> Optional[bool]: - secret_value = get_secret(secret_name=secret_name) + secret_value = _get_env_secret(secret_name=secret_name) if secret_value is None: return default_value @@ -82,70 +81,6 @@ def get_secret_bool( return _str_to_bool(secret_value) -def _get_proxy_signing_key() -> Optional[str]: - salt_key = os.getenv("LITELLM_SALT_KEY") - if salt_key is not None: - return salt_key - - proxy_server_module = sys.modules.get("litellm.proxy.proxy_server") - if proxy_server_module is not None: - proxy_master_key = getattr(proxy_server_module, "master_key", None) - if isinstance(proxy_master_key, str): - return proxy_master_key - - return os.getenv("LITELLM_MASTER_KEY") - - -def _decrypt_value(value: bytes, signing_key: str) -> str: - import nacl.secret - - hash_bytes = hashlib.sha256(signing_key.encode()).digest() - box = nacl.secret.SecretBox(hash_bytes) - - if len(value) == 0: - return "" - - plaintext = box.decrypt(value) - return plaintext.decode("utf-8") - - -def decrypt_value_helper( - value: Any, - key: str, - exception_type: Literal["debug", "error"] = "error", - return_original_value: bool = False, -) -> Any: - signing_key = _get_proxy_signing_key() - - try: - if isinstance(value, str): - if signing_key is None: - raise ValueError("No signing key configured") - - try: - decoded_b64 = base64.urlsafe_b64decode(value) - except Exception: - decoded_b64 = base64.b64decode(value) - - return _decrypt_value(value=decoded_b64, signing_key=signing_key) - - 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 - - verbose_proxy_logger.exception(error_message) - return None - - def _resolve_os_environ_variables(params: dict) -> dict: """ Resolve ``os.environ/`` environment variables in ``litellm_params``. diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index f9127210a3b..5af4d9a5eee 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -17,6 +17,7 @@ import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_mod from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, + _resolve_os_environ_variables, get_callback_identifier, health_license_endpoint, health_services_endpoint, @@ -338,6 +339,70 @@ async def test_health_services_endpoint_sqs(status, error_message): mock_instance.async_health_check.assert_awaited_once() +def test_resolve_os_environ_variables_should_use_secret_manager_get_secret(): + params = { + "api_key": "os.environ/TEST_API_KEY", + "api_base": "https://example.com", + } + + with patch( + "litellm.proxy.health_endpoints._health_endpoints.get_secret", + return_value="resolved-secret-value", + ) as mock_get_secret: + result = _resolve_os_environ_variables(params) + + assert result == { + "api_key": "resolved-secret-value", + "api_base": "https://example.com", + } + mock_get_secret.assert_called_once_with("os.environ/TEST_API_KEY") + + +def test_resolve_os_environ_variables_should_resolve_nested_dicts_and_lists(): + params = { + "api_key": "os.environ/ROOT_SECRET", + "headers": { + "Authorization": "os.environ/AUTH_SECRET", + "static": "value", + }, + "fallbacks": [ + "os.environ/FALLBACK_SECRET", + { + "nested_list_key": "os.environ/NESTED_LIST_SECRET", + }, + ["os.environ/DEEP_LIST_SECRET", "plain-value"], + ], + } + + resolved_values = { + "os.environ/ROOT_SECRET": "root-secret", + "os.environ/AUTH_SECRET": "auth-secret", + "os.environ/FALLBACK_SECRET": "fallback-secret", + "os.environ/NESTED_LIST_SECRET": "nested-list-secret", + "os.environ/DEEP_LIST_SECRET": "deep-list-secret", + } + + with patch( + "litellm.proxy.health_endpoints._health_endpoints.get_secret", + side_effect=lambda secret_name: resolved_values[secret_name], + ) as mock_get_secret: + result = _resolve_os_environ_variables(params) + + assert result == { + "api_key": "root-secret", + "headers": { + "Authorization": "auth-secret", + "static": "value", + }, + "fallbacks": [ + "fallback-secret", + {"nested_list_key": "nested-list-secret"}, + ["deep-list-secret", "plain-value"], + ], + } + assert mock_get_secret.call_count == 5 + + @pytest.mark.asyncio async def test_health_services_endpoint_email_should_use_test_email_address_from_db_when_store_model_in_db_enabled(): mock_prisma = MagicMock()