mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
♻️ refactor(encrypt_decrypt): extract signing key logic to dedicated module
- move _get_salt_key logic to new signing_key_utils module
- use sys.modules lookup to avoid circular import with proxy_server
- add fallback to LITELLM_MASTER_KEY environment variable
✅ test(signing_key_utils): add unit tests for signing key resolution
- test priority: LITELLM_SALT_KEY > proxy_server.master_key > LITELLM_MASTER_KEY
- mock sys.modules to simulate loaded proxy_server module
This commit is contained in:
parent
d34182833c
commit
c87f7c8fcd
3 changed files with 55 additions and 9 deletions
|
|
@ -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):
|
||||
|
|
|
|||
17
litellm/proxy/common_utils/signing_key_utils.py
Normal file
17
litellm/proxy/common_utils/signing_key_utils.py
Normal file
|
|
@ -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")
|
||||
|
|
@ -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"
|
||||
Loading…
Add table
Reference in a new issue