From 8f0b55879f2ae24adef08262dbc4f97ff6983ac7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 6 Jul 2024 11:09:47 -0700 Subject: [PATCH] move encrypt / decrypt to helper --- litellm/proxy/proxy_server.py | 76 ++++++++++------------------------- litellm/proxy/utils.py | 41 ------------------- 2 files changed, 21 insertions(+), 96 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d633075b77..515708df82 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -141,6 +141,10 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth ## Import All Misc routes here ## 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.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.init_callbacks import initialize_callbacks_on_proxy from litellm.proxy.common_utils.openai_endpoint_utils import ( @@ -186,8 +190,6 @@ from litellm.proxy.utils import ( _get_projected_spend_over_limit, _is_projected_spend_over_limit, _is_valid_team_configs, - decrypt_value, - encrypt_value, get_error_message_str, get_instance_fn, hash_token, @@ -1880,16 +1882,8 @@ class ProxyConfig: # decrypt values for k, v in _litellm_params.items(): 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 - _value = decrypt_value(value=decoded_b64, master_key=master_key) + _value = decrypt_value_helper(value=v) # sanity check if string > size 0 if len(_value) > 0: _litellm_params[k] = _value @@ -1933,13 +1927,8 @@ class ProxyConfig: if isinstance(_litellm_params, dict): # decrypt values for k, v in _litellm_params.items(): - if isinstance(v, str): - # decode base64 - decoded_b64 = base64.b64decode(v) - # decrypt value - _litellm_params[k] = decrypt_value( - value=decoded_b64, master_key=master_key # type: ignore - ) + decrypted_value = decrypt_value_helper(value=v) + _litellm_params[k] = decrypted_value _litellm_params = LiteLLM_Params(**_litellm_params) else: verbose_proxy_logger.error( @@ -1995,10 +1984,8 @@ class ProxyConfig: environment_variables = config_data.get("environment_variables", {}) for k, v in environment_variables.items(): try: - if v is not None: - decoded_b64 = base64.b64decode(v) - value = decrypt_value(value=decoded_b64, master_key=master_key) # type: ignore - os.environ[k] = value + decrypted_value = decrypt_value_helper(value=v) + os.environ[k] = decrypted_value except Exception as e: verbose_proxy_logger.error( "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) _orignal_litellm_model_name = model_params.litellm_params.model for k, v in _litellm_params_dict.items(): - if isinstance(v, str): - encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore - model_params.litellm_params[k] = base64.b64encode( - encrypted_value - ).decode("utf-8") + encrypted_value = encrypt_value_helper(value=v) + model_params.litellm_params[k] = encrypted_value _data: dict = { "model_id": model_params.model_info.id, "model_name": model_params.model_name, @@ -6065,11 +6049,8 @@ async def update_model( ### ENCRYPT PARAMS ### for k, v in _new_litellm_params_dict.items(): - if isinstance(v, str): - encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore - model_params.litellm_params[k] = base64.b64encode( - encrypted_value - ).decode("utf-8") + encrypted_value = encrypt_value_helper(value=v) + model_params.litellm_params[k] = encrypted_value ### MERGE WITH EXISTING DATA ### merged_dictionary = {} @@ -8393,11 +8374,8 @@ async def update_config(config_info: ConfigYAML): # encrypt updated_environment_variables # for k, v in _updated_environment_variables.items(): - if isinstance(v, str): - encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore - _updated_environment_variables[k] = base64.b64encode( - encrypted_value - ).decode("utf-8") + encrypted_value = encrypt_value_helper(value=v) + _updated_environment_variables[k] = encrypted_value _existing_env_variables = config["environment_variables"] @@ -8814,11 +8792,8 @@ async def get_config(): env_vars_dict[_var] = None else: # decode + decrypt the value - decoded_b64 = base64.b64decode(env_variable) - _decrypted_value = decrypt_value( - value=decoded_b64, master_key=master_key - ) - env_vars_dict[_var] = _decrypted_value + decrypted_value = decrypt_value_helper(value=env_variable) + env_vars_dict[_var] = decrypted_value _data_to_return.append({"name": _callback, "variables": env_vars_dict}) elif _callback == "langfuse": @@ -8834,11 +8809,8 @@ async def get_config(): _langfuse_env_vars[_var] = None else: # decode + decrypt the value - decoded_b64 = base64.b64decode(env_variable) - _decrypted_value = decrypt_value( - value=decoded_b64, master_key=master_key - ) - _langfuse_env_vars[_var] = _decrypted_value + decrypted_value = decrypt_value_helper(value=env_variable) + _langfuse_env_vars[_var] = decrypted_value _data_to_return.append( {"name": _callback, "variables": _langfuse_env_vars} @@ -8859,10 +8831,7 @@ async def get_config(): _slack_env_vars[_var] = _value else: # decode + decrypt the value - decoded_b64 = base64.b64decode(env_variable) - _decrypted_value = decrypt_value( - value=decoded_b64, master_key=master_key - ) + _decrypted_value = decrypt_value_helper(value=env_variable) _slack_env_vars[_var] = _decrypted_value _alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types @@ -8898,10 +8867,7 @@ async def get_config(): _email_env_vars[_var] = None else: # decode + decrypt the value - decoded_b64 = base64.b64decode(env_variable) - _decrypted_value = decrypt_value( - value=decoded_b64, master_key=master_key - ) + _decrypted_value = decrypt_value_helper(value=env_variable) _email_env_vars[_var] = _decrypted_value alerting_data.append( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 32b74be7c6..5d68c4d3a2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2705,47 +2705,6 @@ def _is_valid_team_configs(team_id=None, team_config=None, request_data=None): 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 url_to_redirect_to = os.getenv("PROXY_BASE_URL", "") url_to_redirect_to += "/login"