diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index a5da5798f47..7d5f6fec54f 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -1,19 +1,12 @@ import base64 -import os from typing import Literal, Optional from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.signing_key_utils import get_proxy_signing_key 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 + return get_proxy_signing_key() def encrypt_value_helper(value: str, new_encryption_key: Optional[str] = None): diff --git a/litellm/proxy/common_utils/signing_key_utils.py b/litellm/proxy/common_utils/signing_key_utils.py new file mode 100644 index 00000000000..f74498cd6f6 --- /dev/null +++ b/litellm/proxy/common_utils/signing_key_utils.py @@ -0,0 +1,17 @@ +import os +import sys +from typing import Optional + + +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") diff --git a/tests/test_litellm/proxy/common_utils/test_signing_key_utils.py b/tests/test_litellm/proxy/common_utils/test_signing_key_utils.py new file mode 100644 index 00000000000..33c22b66595 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_signing_key_utils.py @@ -0,0 +1,36 @@ +import sys +from types import SimpleNamespace + +from litellm.proxy.common_utils.signing_key_utils import get_proxy_signing_key + + +def test_get_proxy_signing_key_prefers_salt_key(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "salt-key") + monkeypatch.setenv("LITELLM_MASTER_KEY", "env-master-key") + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(master_key="proxy-master-key"), + ) + + assert get_proxy_signing_key() == "salt-key" + + +def test_get_proxy_signing_key_uses_loaded_proxy_server_master_key(monkeypatch): + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + monkeypatch.delenv("LITELLM_MASTER_KEY", raising=False) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(master_key="proxy-master-key"), + ) + + assert get_proxy_signing_key() == "proxy-master-key" + + +def test_get_proxy_signing_key_falls_back_to_env_master_key(monkeypatch): + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + monkeypatch.setenv("LITELLM_MASTER_KEY", "env-master-key") + monkeypatch.delitem(sys.modules, "litellm.proxy.proxy_server", raising=False) + + assert get_proxy_signing_key() == "env-master-key"