phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

move encrypt / decrypt to helper

Browse source
This commit is contained in:
Ishaan Jaff 2024-07-06 11:09:47 -07:00
parent 4e7a99e0f9
commit 8f0b55879f
2 changed files with 21 additions and 96 deletions
Download patch file Download diff file Expand all files Collapse all files

76
litellm/proxy/proxy_server.py
View file

@ -141,6 +141,10 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
## Import All Misc routes here ## ## Import All Misc routes here ##
from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.caching_routes import router as caching_router
from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.init_callbacks import initialize_callbacks_on_proxy from litellm.proxy.common_utils.init_callbacks import initialize_callbacks_on_proxy
from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.common_utils.openai_endpoint_utils import (
@ -186,8 +190,6 @@ from litellm.proxy.utils import (
_get_projected_spend_over_limit, _get_projected_spend_over_limit,
_is_projected_spend_over_limit, _is_projected_spend_over_limit,
_is_valid_team_configs, _is_valid_team_configs,
decrypt_value,
encrypt_value,
get_error_message_str, get_error_message_str,
get_instance_fn, get_instance_fn,
hash_token, hash_token,
@ -1880,16 +1882,8 @@ class ProxyConfig:
# decrypt values # decrypt values
for k, v in _litellm_params.items(): for k, v in _litellm_params.items():
if isinstance(v, str): if isinstance(v, str):
# decode base64
try:
decoded_b64 = base64.b64decode(v)
except Exception as e:
verbose_proxy_logger.error(
"Error decoding value - {}".format(v)
)
continue
# decrypt value # decrypt value
_value = decrypt_value(value=decoded_b64, master_key=master_key) _value = decrypt_value_helper(value=v)
# sanity check if string > size 0 # sanity check if string > size 0
if len(_value) > 0: if len(_value) > 0:
_litellm_params[k] = _value _litellm_params[k] = _value
@ -1933,13 +1927,8 @@ class ProxyConfig:
if isinstance(_litellm_params, dict): if isinstance(_litellm_params, dict):
# decrypt values # decrypt values
for k, v in _litellm_params.items(): for k, v in _litellm_params.items():
if isinstance(v, str): decrypted_value = decrypt_value_helper(value=v)
# decode base64 _litellm_params[k] = decrypted_value
decoded_b64 = base64.b64decode(v)
# decrypt value
_litellm_params[k] = decrypt_value(
value=decoded_b64, master_key=master_key # type: ignore
)
_litellm_params = LiteLLM_Params(**_litellm_params) _litellm_params = LiteLLM_Params(**_litellm_params)
else: else:
verbose_proxy_logger.error( verbose_proxy_logger.error(
@ -1995,10 +1984,8 @@ class ProxyConfig:
environment_variables = config_data.get("environment_variables", {}) environment_variables = config_data.get("environment_variables", {})
for k, v in environment_variables.items(): for k, v in environment_variables.items():
try: try:
if v is not None: decrypted_value = decrypt_value_helper(value=v)
decoded_b64 = base64.b64decode(v) os.environ[k] = decrypted_value
value = decrypt_value(value=decoded_b64, master_key=master_key) # type: ignore
os.environ[k] = value
except Exception as e: except Exception as e:
verbose_proxy_logger.error( verbose_proxy_logger.error(
"Error setting env variable: %s - %s", k, str(e) "Error setting env variable: %s - %s", k, str(e)
@ -5930,11 +5917,8 @@ async def add_new_model(
_litellm_params_dict = model_params.litellm_params.dict(exclude_none=True) _litellm_params_dict = model_params.litellm_params.dict(exclude_none=True)
_orignal_litellm_model_name = model_params.litellm_params.model _orignal_litellm_model_name = model_params.litellm_params.model
for k, v in _litellm_params_dict.items(): for k, v in _litellm_params_dict.items():
if isinstance(v, str): encrypted_value = encrypt_value_helper(value=v)
encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore model_params.litellm_params[k] = encrypted_value
model_params.litellm_params[k] = base64.b64encode(
encrypted_value
).decode("utf-8")
_data: dict = { _data: dict = {
"model_id": model_params.model_info.id, "model_id": model_params.model_info.id,
"model_name": model_params.model_name, "model_name": model_params.model_name,
@ -6065,11 +6049,8 @@ async def update_model(
### ENCRYPT PARAMS ### ### ENCRYPT PARAMS ###
for k, v in _new_litellm_params_dict.items(): for k, v in _new_litellm_params_dict.items():
if isinstance(v, str): encrypted_value = encrypt_value_helper(value=v)
encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore model_params.litellm_params[k] = encrypted_value
model_params.litellm_params[k] = base64.b64encode(
encrypted_value
).decode("utf-8")
### MERGE WITH EXISTING DATA ### ### MERGE WITH EXISTING DATA ###
merged_dictionary = {} merged_dictionary = {}
@ -8393,11 +8374,8 @@ async def update_config(config_info: ConfigYAML):
# encrypt updated_environment_variables # # encrypt updated_environment_variables #
for k, v in _updated_environment_variables.items(): for k, v in _updated_environment_variables.items():
if isinstance(v, str): encrypted_value = encrypt_value_helper(value=v)
encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore _updated_environment_variables[k] = encrypted_value
_updated_environment_variables[k] = base64.b64encode(
encrypted_value
).decode("utf-8")
_existing_env_variables = config["environment_variables"] _existing_env_variables = config["environment_variables"]
@ -8814,11 +8792,8 @@ async def get_config():
env_vars_dict[_var] = None env_vars_dict[_var] = None
else: else:
# decode + decrypt the value # decode + decrypt the value
decoded_b64 = base64.b64decode(env_variable) decrypted_value = decrypt_value_helper(value=env_variable)
_decrypted_value = decrypt_value( env_vars_dict[_var] = decrypted_value
value=decoded_b64, master_key=master_key
)
env_vars_dict[_var] = _decrypted_value
_data_to_return.append({"name": _callback, "variables": env_vars_dict}) _data_to_return.append({"name": _callback, "variables": env_vars_dict})
elif _callback == "langfuse": elif _callback == "langfuse":
@ -8834,11 +8809,8 @@ async def get_config():
_langfuse_env_vars[_var] = None _langfuse_env_vars[_var] = None
else: else:
# decode + decrypt the value # decode + decrypt the value
decoded_b64 = base64.b64decode(env_variable) decrypted_value = decrypt_value_helper(value=env_variable)
_decrypted_value = decrypt_value( _langfuse_env_vars[_var] = decrypted_value
value=decoded_b64, master_key=master_key
)
_langfuse_env_vars[_var] = _decrypted_value
_data_to_return.append( _data_to_return.append(
{"name": _callback, "variables": _langfuse_env_vars} {"name": _callback, "variables": _langfuse_env_vars}
@ -8859,10 +8831,7 @@ async def get_config():
_slack_env_vars[_var] = _value _slack_env_vars[_var] = _value
else: else:
# decode + decrypt the value # decode + decrypt the value
decoded_b64 = base64.b64decode(env_variable) _decrypted_value = decrypt_value_helper(value=env_variable)
_decrypted_value = decrypt_value(
value=decoded_b64, master_key=master_key
)
_slack_env_vars[_var] = _decrypted_value _slack_env_vars[_var] = _decrypted_value
_alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types _alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types
@ -8898,10 +8867,7 @@ async def get_config():
_email_env_vars[_var] = None _email_env_vars[_var] = None
else: else:
# decode + decrypt the value # decode + decrypt the value
decoded_b64 = base64.b64decode(env_variable) _decrypted_value = decrypt_value_helper(value=env_variable)
_decrypted_value = decrypt_value(
value=decoded_b64, master_key=master_key
)
_email_env_vars[_var] = _decrypted_value _email_env_vars[_var] = _decrypted_value
alerting_data.append( alerting_data.append(

41
litellm/proxy/utils.py
View file

@ -2705,47 +2705,6 @@ def _is_valid_team_configs(team_id=None, team_config=None, request_data=None):
return return
def encrypt_value(value: str, master_key: str):
import hashlib
import nacl.secret
import nacl.utils
# get 32 byte master key #
hash_object = hashlib.sha256(master_key.encode())
hash_bytes = hash_object.digest()
# initialize secret box #
box = nacl.secret.SecretBox(hash_bytes)
# encode message #
value_bytes = value.encode("utf-8")
encrypted = box.encrypt(value_bytes)
return encrypted
def decrypt_value(value: bytes, master_key: str) -> str:
import hashlib
import nacl.secret
import nacl.utils
# get 32 byte master key #
hash_object = hashlib.sha256(master_key.encode())
hash_bytes = hash_object.digest()
# initialize secret box #
box = nacl.secret.SecretBox(hash_bytes)
# Convert the bytes object to a string
plaintext = box.decrypt(value)
plaintext = plaintext.decode("utf-8") # type: ignore
return plaintext # type: ignore
# LiteLLM Admin UI - Non SSO Login # LiteLLM Admin UI - Non SSO Login
url_to_redirect_to = os.getenv("PROXY_BASE_URL", "") url_to_redirect_to = os.getenv("PROXY_BASE_URL", "")
url_to_redirect_to += "/login" url_to_redirect_to += "/login"