From 466e7d178c10fb1ba018e67e6d287bda3a0efb5d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 5 Nov 2025 14:03:43 -0800 Subject: [PATCH] [Feat] Cyber Ark - Add Key Rotations support (#16289) * KeyManagementSystem add cyberark * add CyberArkSecretManager * add CyberArkSecretManager * add CyberArkSecretManager * docs add CyberArkSecretManager * docs * refactor to use get_secret_from_manager * fix async roate for cyber ark, re-use base class * fixes * cyber ark * docs fix * docs fix * docs cyberark * fix linting * fix get_secret_from_manager --- docs/my-website/docs/proxy/cost_tracking.md | 2 +- .../docs/secret_managers/cyberark.md | 54 +++++--- litellm/__init__.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 1 - .../cyberark_secret_manager.py | 19 +-- litellm/secret_managers/main.py | 5 +- .../secret_managers/secret_manager_handler.py | 2 +- tests/litellm_utils_tests/test_cyberark.py | 120 +++++++++++++++++- 8 files changed, 160 insertions(+), 45 deletions(-) diff --git a/docs/my-website/docs/proxy/cost_tracking.md b/docs/my-website/docs/proxy/cost_tracking.md index da8b6f5c52..019cd62c62 100644 --- a/docs/my-website/docs/proxy/cost_tracking.md +++ b/docs/my-website/docs/proxy/cost_tracking.md @@ -9,7 +9,7 @@ Track spend for keys, users, and teams across 100+ LLMs. LiteLLM automatically tracks spend for all known models. See our [model cost map](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json) :::tip Keep Pricing Data Updated -[Sync model pricing data from GitHub](../sync_models_github.md) to ensure accurate cost tracking. +[Sync model pricing data from GitHub](./sync_models_github.md) to ensure accurate cost tracking. ::: ### How to Track Spend with LiteLLM diff --git a/docs/my-website/docs/secret_managers/cyberark.md b/docs/my-website/docs/secret_managers/cyberark.md index e226fec814..37aa108669 100644 --- a/docs/my-website/docs/secret_managers/cyberark.md +++ b/docs/my-website/docs/secret_managers/cyberark.md @@ -27,7 +27,7 @@ LiteLLM supports two methods of authentication: 1. API key authentication - `CYBERARK_API_KEY` (recommended) 2. Certificate authentication - `CYBERARK_CLIENT_CERT` and `CYBERARK_CLIENT_KEY` -```bash +```bash title="Environment Variables" showLineNumbers CYBERARK_API_BASE="http://your-conjur-instance:8080" CYBERARK_ACCOUNT="default" CYBERARK_USERNAME="admin" @@ -45,7 +45,7 @@ CYBERARK_REFRESH_INTERVAL="300" # defaults to 300 seconds (5 minutes), frequency **Step 2.** Add to proxy config.yaml -```yaml +```yaml title="Proxy Config" showLineNumbers general_settings: key_management_system: "cyberark" @@ -58,7 +58,7 @@ general_settings: **Step 3.** Start + test proxy -```bash +```bash title="Start Proxy" showLineNumbers $ litellm --config /path/to/config.yaml ``` @@ -72,13 +72,13 @@ When you create a virtual key in the LiteLLM UI, it automatically gets stored in In this example, we create a key named `litellm-cyber-ark-secret-key`: -Creating virtual key in LiteLLM UI +Creating virtual key in LiteLLM UI **Step 2:** Verify the secret exists in CyberArk You can verify the virtual key was stored in CyberArk by querying the secrets API: -```bash +```bash title="Verify Secret in CyberArk" showLineNumbers TOKEN=$(curl -s -X POST http://0.0.0.0:8080/authn/default/admin/authenticate \ -d "your-api-key" | base64 | tr -d '\n') @@ -88,7 +88,7 @@ curl -H "Authorization: Token token=\"$TOKEN\"" \ The response shows `litellm-cyber-ark-secret-key` exists in CyberArk: -Virtual key stored in CyberArk API +Virtual key stored in CyberArk API The virtual key is stored with the full path: `default:variable:litellm/litellm-cyber-ark-secret-key` @@ -131,33 +131,49 @@ LiteLLM stores secrets under the `prefix_for_stored_virtual_keys` path (default: For example, a virtual key would be stored as: `litellm/virtual-key-name` -**Working curl examples** +**Important Notes** -Authenticate and get a token: -```bash +- Variables must be defined in a Conjur policy before setting their values +- LiteLLM automatically creates policy entries when writing new secrets +- Secret names with slashes (e.g., `litellm/key`) are automatically URL-encoded +- Session tokens are cached for 5 minutes by default to minimize API calls + +## Troubleshooting + +If you're experiencing issues with the LiteLLM integration, first validate that your CyberArk Conjur instance is working correctly. Run these curl commands directly against your CyberArk endpoints to verify connectivity and authentication: + +**Step 1: Authenticate and get a token** + +Replace `http://conjur.example.com:8080` with your `CYBERARK_API_BASE` and use your actual credentials: + +```bash title="Authenticate" showLineNumbers TOKEN=$(curl -s -X POST http://conjur.example.com:8080/authn/default/admin/authenticate \ -d "your-api-key" | base64 | tr -d '\n') ``` -Read a secret: -```bash +**Step 2: Test reading a secret** + +```bash title="Read Secret" showLineNumbers curl -H "Authorization: Token token=\"$TOKEN\"" \ "http://conjur.example.com:8080/secrets/default/variable/test-secret" ``` -Write a secret: -```bash +**Step 3: Test writing a secret** + +```bash title="Write Secret" showLineNumbers curl -X POST \ -H "Authorization: Token token=\"$TOKEN\"" \ --data "my-secret-value" \ "http://conjur.example.com:8080/secrets/default/variable/test-secret" ``` -**Important Notes** +If these commands work successfully against your CyberArk instance, then CyberArk is functioning correctly and the issue is with your LiteLLM configuration. Check that: +- Your environment variables are correctly set +- The `CYBERARK_API_BASE` URL is accessible from your LiteLLM instance +- Your API key or certificates have the necessary permissions in CyberArk -- Variables must be defined in a Conjur policy before setting their values -- LiteLLM automatically creates policy entries when writing new secrets -- Secret names with slashes (e.g., `litellm/key`) are automatically URL-encoded -- CyberArk Conjur does not support direct secret deletion via API (must use policy updates) -- Session tokens are cached for 5 minutes by default to minimize API calls +## Video Walkthrough +This video walks through using CyberArk Conjur as a secret manager with LiteLLM. We create a virtual key in the LiteLLM Admin UI and verify it exists in CyberArk. Then we rotate the secret key and verify it exists in CyberArk. + + diff --git a/litellm/__init__.py b/litellm/__init__.py index c253b6b516..29d83415e7 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -5,7 +5,7 @@ warnings.filterwarnings("ignore", message=".*conflict with protected namespace.* # Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances # This warning can accumulate during streaming and cause memory leaks warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*") -### INIT VARIABLES ###################### +### INIT VARIABLES ####################### import threading import os from typing import ( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 458b02cc0d..d7be4d296b 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1,7 +1,6 @@ # What is this? ## Common Utility file for Logging handler # Logging function -> log the exact model details + what's being sent | Non-Blocking -import asyncio import copy import datetime import json diff --git a/litellm/secret_managers/cyberark_secret_manager.py b/litellm/secret_managers/cyberark_secret_manager.py index 5aad215453..0d439290c7 100644 --- a/litellm/secret_managers/cyberark_secret_manager.py +++ b/litellm/secret_managers/cyberark_secret_manager.py @@ -127,12 +127,12 @@ class CyberArkSecretManager(BaseSecretManager): content=policy_yaml, ) resp.raise_for_status() - verbose_logger.debug("Created policy entry for a variable.") + verbose_logger.debug(f"Created policy entry for variable: {secret_name}") except httpx.HTTPStatusError as e: # Variable might already exist, which is fine if e.response.status_code in [409, 422]: verbose_logger.debug( - "A variable already exists or policy conflict (expected)" + f"Variable {secret_name} already exists or policy conflict (expected)" ) else: verbose_logger.warning( @@ -303,21 +303,6 @@ class CyberArkSecretManager(BaseSecretManager): verbose_logger.exception(f"Error writing secret to CyberArk Conjur: {e}") return {"status": "error", "message": str(e)} - async def async_rotate_secret( - self, - current_secret_name: str, - new_secret_name: str, - new_secret_value: str, - optional_params: Optional[Dict] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Dict: - """ - CyberArk Conjur does not have built-in secret rotation. - - Raises: - NotImplementedError: Always raised - """ - raise NotImplementedError("CyberArk Conjur does not support secret rotation") async def async_delete_secret( self, diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 51a7f96ecb..a093fe2d2f 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -1,19 +1,18 @@ import ast import os import traceback -from typing import Any, Optional, Union +from typing import Optional, Union import httpx import litellm -from litellm._logging import print_verbose, verbose_logger +from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, ) from litellm.secret_managers.secret_manager_handler import get_secret_from_manager -from litellm.types.secret_managers.main import KeyManagementSystem oidc_cache = DualCache() diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index c150c6ac1d..fd991296ef 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -21,7 +21,7 @@ def _is_base64(s): return False -def get_secret_from_manager( +def get_secret_from_manager( # noqa: PLR0915 client: Any, key_manager: str, secret_name: str, diff --git a/tests/litellm_utils_tests/test_cyberark.py b/tests/litellm_utils_tests/test_cyberark.py index d77ec85e0f..45783b99fe 100644 --- a/tests/litellm_utils_tests/test_cyberark.py +++ b/tests/litellm_utils_tests/test_cyberark.py @@ -39,16 +39,132 @@ async def test_cyberark_write_and_read_secret(): secret_name=secret_name, secret_value=secret_value, ) - # Avoid logging write_response to prevent leaking secret names + print("write_response=", write_response) # Validate write was successful assert write_response["status"] == "success" # Read the secret back read_value = cyberark_manager.sync_read_secret(secret_name=secret_name) - # Don't log secret value in clear text + print("READ VALUE=", read_value) # Validate the secret exists and has the correct value assert read_value is not None assert read_value == secret_value + +@pytest.mark.asyncio +async def test_cyberark_rotate_secret(): + """ + Integration test: Test key rotation in CyberArk Conjur. + + This test simulates what happens when a virtual key is rotated: + 1. Write initial secret with alias (like sk-1234) + 2. Rotate to new value (like sk-12359) + 3. Verify reading the secret returns the NEW value + """ + with patch("litellm.proxy.proxy_server.premium_user", True): + # Create CyberArk secret manager instance + cyberark_manager = CyberArkSecretManager() + + # Simulate initial virtual key creation + secret_alias = f"test-rotation-key-{uuid.uuid4()}" + initial_key_value = f"sk-initial-{uuid.uuid4()}" + rotated_key_value = f"sk-rotated-{uuid.uuid4()}" + + print(f"\n=== Testing Key Rotation ===") + print(f"Alias: {secret_alias}") + print(f"Initial value: {initial_key_value}") + print(f"Rotated value: {rotated_key_value}") + + # Step 1: Write initial secret (simulates key creation) + write_response = await cyberark_manager.async_write_secret( + secret_name=secret_alias, + secret_value=initial_key_value, + ) + print(f"\n1. Initial write response: {write_response}") + assert write_response["status"] == "success" + + # Verify initial value was written + initial_read = cyberark_manager.sync_read_secret(secret_name=secret_alias) + print(f"2. Initial read value: {initial_read}") + assert initial_read == initial_key_value + + # Step 2: Rotate the secret (simulates key rotation) + # In key rotation, we keep the same secret_name but update the value + rotation_response = await cyberark_manager.async_rotate_secret( + current_secret_name=secret_alias, + new_secret_name=secret_alias, # Same name = update in place + new_secret_value=rotated_key_value, + ) + print(f"3. Rotation response: {rotation_response}") + assert rotation_response["status"] == "success" + + # Step 3: Verify the secret now returns the NEW value + rotated_read = cyberark_manager.sync_read_secret(secret_name=secret_alias) + print(f"4. After rotation, read value: {rotated_read}") + + # This is the key assertion: after rotation, reading should return the NEW value + assert rotated_read is not None + assert rotated_read == rotated_key_value + assert rotated_read != initial_key_value + + print(f"\nāœ… Rotation successful: {initial_key_value} ā†’ {rotated_key_value}") + + +@pytest.mark.asyncio +async def test_cyberark_rotate_secret_with_new_alias(): + """ + Integration test: Test key rotation with a new alias. + + This simulates rotating a key and changing its alias at the same time: + 1. Write secret with alias-v1 + 2. Rotate to alias-v2 with new value + 3. Verify alias-v2 has the new value + 4. Verify alias-v1 still exists with old value (CyberArk doesn't delete) + """ + with patch("litellm.proxy.proxy_server.premium_user", True): + # Create CyberArk secret manager instance + cyberark_manager = CyberArkSecretManager() + + # Simulate key rotation with alias change + base_alias = f"test-alias-change-{uuid.uuid4()}" + old_alias = f"{base_alias}-v1" + new_alias = f"{base_alias}-v2" + old_value = f"sk-old-{uuid.uuid4()}" + new_value = f"sk-new-{uuid.uuid4()}" + + print(f"\n=== Testing Key Rotation with Alias Change ===") + print(f"Old alias: {old_alias} = {old_value}") + print(f"New alias: {new_alias} = {new_value}") + + # Step 1: Create initial secret with old alias + write_response = await cyberark_manager.async_write_secret( + secret_name=old_alias, + secret_value=old_value, + ) + print(f"\n1. Initial write: {write_response}") + assert write_response["status"] == "success" + + # Step 2: Rotate to new alias with new value + rotation_response = await cyberark_manager.async_rotate_secret( + current_secret_name=old_alias, + new_secret_name=new_alias, # Different name = new secret + new_secret_value=new_value, + ) + print(f"2. Rotation response: {rotation_response}") + assert rotation_response["status"] == "success" + + # Step 3: Verify new alias has new value + new_read = cyberark_manager.sync_read_secret(secret_name=new_alias) + print(f"3. Read new alias: {new_read}") + assert new_read == new_value + + # Step 4: Verify old alias still exists (CyberArk doesn't delete via API) + old_read = cyberark_manager.sync_read_secret(secret_name=old_alias) + print(f"4. Read old alias (should still exist): {old_read}") + assert old_read == old_value # Old secret still exists in CyberArk + + print(f"\nāœ… Alias rotation successful: {old_alias} ā†’ {new_alias}") + print(f" Note: Old alias still exists in CyberArk (expected behavior)") +