mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
♻️ refactor(health_endpoints): extract encryption utilities and update secret resolution
- 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
This commit is contained in:
yangdx
parent
e2a3ccccc4
commit
d34182833c
2 changed files with 69 additions and 69 deletions
|
|
@ -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``.
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue