From f24e9865341779b07b252f26d7cd97d954017222 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Fri, 16 Aug 2024 20:11:24 +0000 Subject: [PATCH 1/2] (oidc): Add support for loading tokens via a file, environment variable, and from a file path set in an env var. --- litellm/tests/test_secret_manager.py | 61 ++++++++++++++++++++++++++++ litellm/utils.py | 19 +++++++++ 2 files changed, 80 insertions(+) diff --git a/litellm/tests/test_secret_manager.py b/litellm/tests/test_secret_manager.py index 904b291bfe..652e209895 100644 --- a/litellm/tests/test_secret_manager.py +++ b/litellm/tests/test_secret_manager.py @@ -5,6 +5,8 @@ from dotenv import load_dotenv load_dotenv() import os +from uuid import uuid4 +import tempfile sys.path.insert( 0, os.path.abspath("../..") @@ -135,3 +137,62 @@ def test_oidc_circle_v1_with_amazon_fips(): aws_session_name="assume-v1-session-fips", aws_sts_endpoint="https://sts-fips.us-west-1.amazonaws.com", ) + + +def test_oidc_env_variable(): + # Create a unique environment variable name + env_var_name = "OIDC_TEST_PATH_" + uuid4().hex + os.environ[env_var_name] = "secret-" + uuid4().hex + secret_val = get_secret( + f"oidc/env/{env_var_name}" + ) + + print(f"secret_val: {redact_oidc_signature(secret_val)}") + + assert secret_val == os.environ[env_var_name] + + # now unset the environment variable + del os.environ[env_var_name] + + +def test_oidc_file(): + # Create a temporary file + with tempfile.NamedTemporaryFile(mode='w+') as temp_file: + secret_value = "secret-" + uuid4().hex + temp_file.write(secret_value) + temp_file.flush() + temp_file_path = temp_file.name + + secret_val = get_secret( + f"oidc/file/{temp_file_path}" + ) + + print(f"secret_val: {redact_oidc_signature(secret_val)}") + + assert secret_val == secret_value + + +def test_oidc_env_path(): + # Create a temporary file + with tempfile.NamedTemporaryFile(mode='w+') as temp_file: + secret_value = "secret-" + uuid4().hex + temp_file.write(secret_value) + temp_file.flush() + temp_file_path = temp_file.name + + # Create a unique environment variable name + env_var_name = "OIDC_TEST_PATH_" + uuid4().hex + + # Set the environment variable to the temporary file path + os.environ[env_var_name] = temp_file_path + + # Test getting the secret using the environment variable + secret_val = get_secret( + f"oidc/env_path/{env_var_name}" + ) + + print(f"secret_val: {redact_oidc_signature(secret_val)}") + + assert secret_val == secret_value + + del os.environ[env_var_name] diff --git a/litellm/utils.py b/litellm/utils.py index 40564c1077..4746e7f845 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8433,6 +8433,25 @@ def get_secret( with open(azure_federated_token_file, "r") as f: oidc_token = f.read() return oidc_token + elif oidc_provider == "file": + # Load token from a file + with open(oidc_aud, "r") as f: + oidc_token = f.read() + return oidc_token + elif oidc_provider == "env": + # Load token directly from an environment variable + oidc_token = os.getenv(oidc_aud) + if oidc_token is None: + raise ValueError(f"Environment variable {oidc_aud} not found") + return oidc_token + elif oidc_provider == "env_path": + # Load token from a file path specified in an environment variable + token_file_path = os.getenv(oidc_aud) + if token_file_path is None: + raise ValueError(f"Environment variable {oidc_aud} not found") + with open(token_file_path, "r") as f: + oidc_token = f.read() + return oidc_token else: raise ValueError("Unsupported OIDC provider") From 3a9125d947993e12995ddaf763a5b7bb7c08835f Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Fri, 16 Aug 2024 20:30:41 +0000 Subject: [PATCH 2/2] (oidc): Improve docs for unofficial provider. --- docs/my-website/docs/oidc.md | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/docs/my-website/docs/oidc.md b/docs/my-website/docs/oidc.md index 1f792f8d23..f30edf5044 100644 --- a/docs/my-website/docs/oidc.md +++ b/docs/my-website/docs/oidc.md @@ -19,9 +19,17 @@ LiteLLM supports the following OIDC identity providers: | CircleCI v2 | `circleci_v2`| No | | GitHub Actions | `github` | Yes | | Azure Kubernetes Service | `azure` | No | +| File | `file` | No | +| Environment Variable | `env` | No | +| Environment Path | `env_path` | No | If you would like to use a different OIDC provider, please open an issue on GitHub. +:::tip + +Do not use the `file`, `env`, or `env_path` providers unless you know what you're doing, and you are sure none of the other providers will work for your use-case. Hint: they probably will. + +::: ## OIDC Connect Relying Party (RP) @@ -46,6 +54,32 @@ For providers that do not use the `audience` parameter, you can (and should) omi oidc/config_name_here/ ``` +#### Unofficial Providers (not recommended) + +For the unofficial `file` provider, you can use the following format: + +``` +oidc/file/home/user/dave/this_is_a_file_with_a_token.txt +``` + +For the unofficial `env`, use the following format, where `SECRET_TOKEN` is the name of the environment variable that contains the token: + +``` +oidc/env/SECRET_TOKEN +``` + +For the unofficial `env_path`, use the following format, where `SECRET_TOKEN` is the name of the environment variable that contains the path to the file with the token: + +``` +oidc/env_path/SECRET_TOKEN +``` + +:::tip + +If you are tempted to use oidc/env_path/AZURE_FEDERATED_TOKEN_FILE, don't do that. Instead, use `oidc/azure/`, as this will ensure continued support from LiteLLM if Azure changes their OIDC configuration and/or adds new features. + +::: + ## Examples ### Google Cloud Run -> Amazon Bedrock