From bbf3de74dd260ecd09bac8adcc2dd115cbf067cb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 8 Jul 2024 07:43:00 -0700 Subject: [PATCH 1/2] fix - setting rpm/tpm --- litellm/proxy/common_utils/encrypt_decrypt_utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index f0090046b5..09db1c6bf3 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -23,9 +23,11 @@ def encrypt_value_helper(value: str): return encrypted_value - raise ValueError( + verbose_proxy_logger.debug( f"Invalid value type passed to encrypt_value: {type(value)} for Value: {value}\n Value must be a string" ) + # if it's not a string - do not encrypt it and return the value + return value except Exception as e: raise e @@ -42,6 +44,9 @@ def decrypt_value_helper(value: str): decoded_b64 = base64.b64decode(value) value = decrypt_value(value=decoded_b64, signing_key=signing_key) # type: ignore return value + + # if it's not str - do not decrypt it, return the value + return value except Exception as e: verbose_proxy_logger.error(f"Error decrypting value: {value}\nError: {str(e)}") # [Non-Blocking Exception. - this should not block decrypting other values] From db1d9d4799bcf51c90ec7aa0e33daf96ef464ee6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 8 Jul 2024 07:54:07 -0700 Subject: [PATCH 2/2] test test_encrypt_decrypt --- litellm/tests/test_proxy_encrypt_decrypt.py | 45 +++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 litellm/tests/test_proxy_encrypt_decrypt.py diff --git a/litellm/tests/test_proxy_encrypt_decrypt.py b/litellm/tests/test_proxy_encrypt_decrypt.py new file mode 100644 index 0000000000..6db37e4c13 --- /dev/null +++ b/litellm/tests/test_proxy_encrypt_decrypt.py @@ -0,0 +1,45 @@ +import os +import sys + +import pytest +from dotenv import load_dotenv + +load_dotenv() +import io +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds-the parent directory to the system path + +from litellm.proxy import proxy_server +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) + + +def test_encrypt_decrypt_with_master_key(): + setattr(proxy_server, "master_key", "sk-1234") + assert decrypt_value_helper(encrypt_value_helper("test")) == "test" + assert decrypt_value_helper(encrypt_value_helper(10)) == 10 + assert decrypt_value_helper(encrypt_value_helper(True)) is True + assert decrypt_value_helper(encrypt_value_helper(None)) is None + assert decrypt_value_helper(encrypt_value_helper({"rpm": 10})) == {"rpm": 10} + + # encryption should actually occur for strings + assert encrypt_value_helper("test") != "test" + + +def test_encrypt_decrypt_with_salt_key(): + os.environ["LITELLM_SALT_KEY"] = "sk-salt-key2222" + assert decrypt_value_helper(encrypt_value_helper("test")) == "test" + assert decrypt_value_helper(encrypt_value_helper(10)) == 10 + assert decrypt_value_helper(encrypt_value_helper(True)) is True + assert decrypt_value_helper(encrypt_value_helper(None)) is None + assert decrypt_value_helper(encrypt_value_helper({"rpm": 10})) == {"rpm": 10} + + # encryption should actually occur for strings + assert encrypt_value_helper("test") != "test" + + os.environ.pop("LITELLM_SALT_KEY", None)