From 11668c31c134df6c11683d39d53f2bb7aee5cc55 Mon Sep 17 00:00:00 2001
From: David Manouchehri <david.manouchehri@ai.moda>
Date: Fri, 16 Aug 2024 20:11:24 +0000
Subject: [PATCH] (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 904b291bf..652e20989 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 40564c107..4746e7f84 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")