mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 02:34:29 +00:00
(Feat) Hashicorp Secret Manager - Allow storing virtual keys in secret manager (#7549)
All checks were successful
Read Version from pyproject.toml / read-version (push) Successful in 13s
All checks were successful
Read Version from pyproject.toml / read-version (push) Successful in 13s
* use a base abstract class * async_write_secret for hcorp * fix hcorp * async_write_secret for hashicopr secret manager * store virtual keys in hcorp * add delete secret * test_hashicorp_secret_manager_write_secret * test_hashicorp_secret_manager_delete_secret * docs Supported Secret Managers * docs storing keys in hcorp * docs hcorp * docs secret managers * test_key_generate_with_secret_manager_call * fix unused imports
This commit is contained in:
parent
7f7222ce30
commit
46d9d29bff
13 changed files with 458 additions and 119 deletions
|
@ -21,7 +21,7 @@ This covers:
|
|||
- ✅ [Audit Logs with retention policy](./proxy/enterprise#audit-logs)
|
||||
- ✅ [JWT-Auth](../docs/proxy/token_auth.md)
|
||||
- ✅ [Control available public, private routes (Restrict certain endpoints on proxy)](./proxy/enterprise#control-available-public-private-routes)
|
||||
- ✅ [**Secret Managers** AWS Key Manager, Google Secret Manager, Azure Key](./secret)
|
||||
- ✅ [**Secret Managers** AWS Key Manager, Google Secret Manager, Azure Key, Hashicorp Vault](./secret)
|
||||
- ✅ IP address‑based access control lists
|
||||
- ✅ Track Request IP Address
|
||||
- ✅ [Use LiteLLM keys/authentication on Pass Through Endpoints](./proxy/pass_through#✨-enterprise---use-litellm-keysauthentication-on-pass-through-endpoints)
|
||||
|
|
|
@ -17,6 +17,7 @@ Features:
|
|||
- ✅ [JWT-Auth](../docs/proxy/token_auth.md)
|
||||
- ✅ [Control available public, private routes (Restrict certain endpoints on proxy)](#control-available-public-private-routes)
|
||||
- ✅ [Control available public, private routes](#control-available-public-private-routes)
|
||||
- ✅ [Secret Managers - AWS Key Manager, Google Secret Manager, Azure Key, Hashicorp Vault](../secret)
|
||||
- ✅ [[BETA] AWS Key Manager v2 - Key Decryption](#beta-aws-key-manager---key-decryption)
|
||||
- ✅ IP address‑based access control lists
|
||||
- ✅ Track Request IP Address
|
||||
|
|
|
@ -3,7 +3,6 @@ import TabItem from '@theme/TabItem';
|
|||
import Image from '@theme/IdealImage';
|
||||
|
||||
# Secret Manager
|
||||
LiteLLM supports reading secrets from Azure Key Vault, Google Secret Manager
|
||||
|
||||
:::info
|
||||
|
||||
|
@ -15,6 +14,8 @@ LiteLLM supports reading secrets from Azure Key Vault, Google Secret Manager
|
|||
|
||||
:::
|
||||
|
||||
LiteLLM supports **reading secrets (eg. `OPENAI_API_KEY`)** and **writing secrets (eg. Virtual Keys)** from Azure Key Vault, Google Secret Manager, Hashicorp Vault, and AWS Secret Manager.
|
||||
|
||||
## Supported Secret Managers
|
||||
|
||||
- AWS Key Management Service
|
||||
|
@ -23,37 +24,17 @@ LiteLLM supports reading secrets from Azure Key Vault, Google Secret Manager
|
|||
- [Google Secret Manager](#google-secret-manager)
|
||||
- Google Key Management Service
|
||||
- [Hashicorp Vault](#hashicorp-vault)
|
||||
- [Infisical Secret Manager](#infisical-secret-manager)
|
||||
- [.env Files](#env-files)
|
||||
|
||||
## AWS Key Management V1
|
||||
|
||||
:::tip
|
||||
|
||||
[BETA] AWS Key Management v2 is on the enterprise tier. Go [here for docs](./proxy/enterprise.md#beta-aws-key-manager---key-decryption)
|
||||
|
||||
:::
|
||||
|
||||
Use AWS KMS to storing a hashed copy of your Proxy Master Key in the environment.
|
||||
|
||||
```bash
|
||||
export LITELLM_MASTER_KEY="djZ9xjVaZ..." # 👈 ENCRYPTED KEY
|
||||
export AWS_REGION_NAME="us-west-2"
|
||||
```
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
key_management_system: "aws_kms"
|
||||
key_management_settings:
|
||||
hosted_keys: ["LITELLM_MASTER_KEY"] # 👈 WHICH KEYS ARE STORED ON KMS
|
||||
```
|
||||
|
||||
[**See Decryption Code**](https://github.com/BerriAI/litellm/blob/a2da2a8f168d45648b61279d4795d647d94f90c9/litellm/utils.py#L10182)
|
||||
|
||||
## AWS Secret Manager
|
||||
|
||||
Store your proxy keys in AWS Secret Manager.
|
||||
|
||||
|
||||
| Feature | Support | Description |
|
||||
|---------|----------|-------------|
|
||||
| Reading Secrets | ✅ | Read secrets e.g `OPENAI_API_KEY` |
|
||||
| Writing Secrets | ✅ | Store secrets e.g `Virtual Keys` |
|
||||
|
||||
#### Proxy Usage
|
||||
|
||||
1. Save AWS Credentials in your environment
|
||||
|
@ -100,6 +81,110 @@ general_settings:
|
|||
litellm --config /path/to/config.yaml
|
||||
```
|
||||
|
||||
|
||||
## Hashicorp Vault
|
||||
|
||||
|
||||
| Feature | Support | Description |
|
||||
|---------|----------|-------------|
|
||||
| Reading Secrets | ✅ | Read secrets e.g `OPENAI_API_KEY` |
|
||||
| Writing Secrets | ✅ | Store secrets e.g `Virtual Keys` |
|
||||
|
||||
Read secrets from [Hashicorp Vault](https://developer.hashicorp.com/vault/docs/secrets/kv/kv-v2)
|
||||
|
||||
**Step 1.** Add Hashicorp Vault details in your environment
|
||||
|
||||
LiteLLM supports two methods of authentication:
|
||||
|
||||
1. TLS cert authentication - `HCP_VAULT_CLIENT_CERT` and `HCP_VAULT_CLIENT_KEY`
|
||||
2. Token authentication - `HCP_VAULT_TOKEN`
|
||||
|
||||
```bash
|
||||
HCP_VAULT_ADDR="https://test-cluster-public-vault-0f98180c.e98296b2.z1.hashicorp.cloud:8200"
|
||||
HCP_VAULT_NAMESPACE="admin"
|
||||
|
||||
# Authentication via TLS cert
|
||||
HCP_VAULT_CLIENT_CERT="path/to/client.pem"
|
||||
HCP_VAULT_CLIENT_KEY="path/to/client.key"
|
||||
|
||||
# OR - Authentication via token
|
||||
HCP_VAULT_TOKEN="hvs.CAESIG52gL6ljBSdmq*****"
|
||||
|
||||
|
||||
# OPTIONAL
|
||||
HCP_VAULT_REFRESH_INTERVAL="86400" # defaults to 86400, frequency of cache refresh for Hashicorp Vault
|
||||
```
|
||||
|
||||
**Step 2.** Add to proxy config.yaml
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
key_management_system: "hashicorp_vault"
|
||||
|
||||
# [OPTIONAL SETTINGS]
|
||||
key_management_settings:
|
||||
store_virtual_keys: true # OPTIONAL. Defaults to False, when True will store virtual keys in secret manager
|
||||
prefix_for_stored_virtual_keys: "litellm/" # OPTIONAL. If set, this prefix will be used for stored virtual keys in the secret manager
|
||||
access_mode: "read_and_write" # Literal["read_only", "write_only", "read_and_write"]
|
||||
```
|
||||
|
||||
**Step 3.** Start + test proxy
|
||||
|
||||
```
|
||||
$ litellm --config /path/to/config.yaml
|
||||
```
|
||||
|
||||
[Quick Test Proxy](./proxy/user_keys)
|
||||
|
||||
|
||||
#### How it works
|
||||
|
||||
**Reading Secrets**
|
||||
LiteLLM reads secrets from Hashicorp Vault's KV v2 engine using the following URL format:
|
||||
```
|
||||
{VAULT_ADDR}/v1/{NAMESPACE}/secret/data/{SECRET_NAME}
|
||||
```
|
||||
|
||||
For example, if you have:
|
||||
- `HCP_VAULT_ADDR="https://vault.example.com:8200"`
|
||||
- `HCP_VAULT_NAMESPACE="admin"`
|
||||
- Secret name: `AZURE_API_KEY`
|
||||
|
||||
|
||||
LiteLLM will look up:
|
||||
```
|
||||
https://vault.example.com:8200/v1/admin/secret/data/AZURE_API_KEY
|
||||
```
|
||||
|
||||
#### Expected Secret Format
|
||||
LiteLLM expects all secrets to be stored as a JSON object with a `key` field containing the secret value.
|
||||
|
||||
For example, for `AZURE_API_KEY`, the secret should be stored as:
|
||||
|
||||
```json
|
||||
{
|
||||
"key": "sk-1234"
|
||||
}
|
||||
```
|
||||
|
||||
<Image img={require('../img/hcorp.png')} />
|
||||
|
||||
**Writing Secrets**
|
||||
|
||||
When a Virtual Key is Created / Deleted on LiteLLM, LiteLLM will automatically create / delete the secret in Hashicorp Vault.
|
||||
|
||||
- Create Virtual Key on LiteLLM either through the LiteLLM Admin UI or API
|
||||
|
||||
<Image img={require('../img/hcorp_create_virtual_key.png')} />
|
||||
|
||||
|
||||
- Check Hashicorp Vault for secret
|
||||
|
||||
LiteLLM stores secret under the `prefix_for_stored_virtual_keys` path (default: `litellm/`)
|
||||
|
||||
<Image img={require('../img/hcorp_virtual_key.png')} />
|
||||
|
||||
|
||||
## Azure Key Vault
|
||||
<!--
|
||||
### Quick Start
|
||||
|
@ -240,82 +325,31 @@ $ litellm --test
|
|||
## .env Files
|
||||
If no secret manager client is specified, Litellm automatically uses the `.env` file to manage sensitive data. -->
|
||||
|
||||
## Hashicorp Vault
|
||||
## AWS Key Management V1
|
||||
|
||||
Read secrets from [Hashicorp Vault](https://developer.hashicorp.com/vault/docs/secrets/kv/kv-v2)
|
||||
:::tip
|
||||
|
||||
Step 1. Add Hashicorp Vault details in your environment
|
||||
[BETA] AWS Key Management v2 is on the enterprise tier. Go [here for docs](./proxy/enterprise.md#beta-aws-key-manager---key-decryption)
|
||||
|
||||
LiteLLM supports two methods of authentication:
|
||||
:::
|
||||
|
||||
1. TLS cert authentication - `HCP_VAULT_CLIENT_CERT` and `HCP_VAULT_CLIENT_KEY`
|
||||
2. Token authentication - `HCP_VAULT_TOKEN`
|
||||
Use AWS KMS to storing a hashed copy of your Proxy Master Key in the environment.
|
||||
|
||||
```bash
|
||||
HCP_VAULT_ADDR="https://test-cluster-public-vault-0f98180c.e98296b2.z1.hashicorp.cloud:8200"
|
||||
HCP_VAULT_NAMESPACE="admin"
|
||||
|
||||
# Authentication via TLS cert
|
||||
HCP_VAULT_CLIENT_CERT="path/to/client.pem"
|
||||
HCP_VAULT_CLIENT_KEY="path/to/client.key"
|
||||
|
||||
# OR - Authentication via token
|
||||
HCP_VAULT_TOKEN="hvs.CAESIG52gL6ljBSdmq*****"
|
||||
|
||||
|
||||
# OPTIONAL
|
||||
HCP_VAULT_REFRESH_INTERVAL="86400" # defaults to 86400, frequency of cache refresh for Hashicorp Vault
|
||||
export LITELLM_MASTER_KEY="djZ9xjVaZ..." # 👈 ENCRYPTED KEY
|
||||
export AWS_REGION_NAME="us-west-2"
|
||||
```
|
||||
|
||||
Step 2. Add to proxy config.yaml
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
key_management_system: "hashicorp_vault"
|
||||
key_management_system: "aws_kms"
|
||||
key_management_settings:
|
||||
hosted_keys: ["LITELLM_MASTER_KEY"] # 👈 WHICH KEYS ARE STORED ON KMS
|
||||
```
|
||||
|
||||
Step 3. Start + test proxy
|
||||
[**See Decryption Code**](https://github.com/BerriAI/litellm/blob/a2da2a8f168d45648b61279d4795d647d94f90c9/litellm/utils.py#L10182)
|
||||
|
||||
```
|
||||
$ litellm --config /path/to/config.yaml
|
||||
```
|
||||
|
||||
[Quick Test Proxy](./proxy/user_keys)
|
||||
|
||||
|
||||
#### How it works
|
||||
|
||||
LiteLLM reads secrets from Hashicorp Vault's KV v2 engine using the following URL format:
|
||||
```
|
||||
{VAULT_ADDR}/v1/{NAMESPACE}/secret/data/{SECRET_NAME}
|
||||
```
|
||||
|
||||
For example, if you have:
|
||||
- `HCP_VAULT_ADDR="https://vault.example.com:8200"`
|
||||
- `HCP_VAULT_NAMESPACE="admin"`
|
||||
- Secret name: `AZURE_API_KEY`
|
||||
|
||||
|
||||
LiteLLM will look up:
|
||||
```
|
||||
https://vault.example.com:8200/v1/admin/secret/data/AZURE_API_KEY
|
||||
```
|
||||
|
||||
#### Expected Secret Format
|
||||
LiteLLM expects all secrets to be stored as a JSON object with a `key` field containing the secret value.
|
||||
|
||||
For example, for `AZURE_API_KEY`, the secret should be stored as:
|
||||
|
||||
```json
|
||||
{
|
||||
"key": "sk-1234"
|
||||
}
|
||||
```
|
||||
|
||||
<Image img={require('../img/hcorp.png')} />
|
||||
|
||||
|
||||
## All Secret Manager Settings
|
||||
## **All Secret Manager Settings**
|
||||
|
||||
All settings related to secret management
|
||||
|
||||
|
|
BIN
docs/my-website/img/hcorp_create_virtual_key.png
Normal file
BIN
docs/my-website/img/hcorp_create_virtual_key.png
Normal file
Binary file not shown.
After Width: | Height: | Size: 190 KiB |
BIN
docs/my-website/img/hcorp_virtual_key.png
Normal file
BIN
docs/my-website/img/hcorp_virtual_key.png
Normal file
Binary file not shown.
After Width: | Height: | Size: 145 KiB |
|
@ -10,7 +10,6 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
GenerateKeyRequest,
|
||||
KeyManagementSystem,
|
||||
KeyRequest,
|
||||
LiteLLM_AuditLogs,
|
||||
LiteLLM_VerificationToken,
|
||||
|
@ -195,21 +194,28 @@ class KeyManagementEventHooks:
|
|||
"""
|
||||
if litellm._key_management_settings is not None:
|
||||
if litellm._key_management_settings.store_virtual_keys is True:
|
||||
from litellm.secret_managers.aws_secret_manager_v2 import (
|
||||
AWSSecretsManagerV2,
|
||||
from litellm.secret_managers.base_secret_manager import (
|
||||
BaseSecretManager,
|
||||
)
|
||||
|
||||
# store the key in the secret manager
|
||||
if (
|
||||
litellm._key_management_system
|
||||
== KeyManagementSystem.AWS_SECRET_MANAGER
|
||||
and isinstance(litellm.secret_manager_client, AWSSecretsManagerV2)
|
||||
):
|
||||
if isinstance(litellm.secret_manager_client, BaseSecretManager):
|
||||
await litellm.secret_manager_client.async_write_secret(
|
||||
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{secret_name}",
|
||||
secret_name=KeyManagementEventHooks._get_secret_name(
|
||||
secret_name
|
||||
),
|
||||
secret_value=secret_token,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_secret_name(secret_name: str) -> str:
|
||||
if litellm._key_management_settings.prefix_for_stored_virtual_keys.endswith(
|
||||
"/"
|
||||
):
|
||||
return f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{secret_name}"
|
||||
else:
|
||||
return f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{secret_name}"
|
||||
|
||||
@staticmethod
|
||||
async def _delete_virtual_keys_from_secret_manager(
|
||||
keys_being_deleted: List[LiteLLM_VerificationToken],
|
||||
|
@ -222,15 +228,17 @@ class KeyManagementEventHooks:
|
|||
"""
|
||||
if litellm._key_management_settings is not None:
|
||||
if litellm._key_management_settings.store_virtual_keys is True:
|
||||
from litellm.secret_managers.aws_secret_manager_v2 import (
|
||||
AWSSecretsManagerV2,
|
||||
from litellm.secret_managers.base_secret_manager import (
|
||||
BaseSecretManager,
|
||||
)
|
||||
|
||||
if isinstance(litellm.secret_manager_client, AWSSecretsManagerV2):
|
||||
if isinstance(litellm.secret_manager_client, BaseSecretManager):
|
||||
for key in keys_being_deleted:
|
||||
if key.key_alias is not None:
|
||||
await litellm.secret_manager_client.async_delete_secret(
|
||||
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key.key_alias}"
|
||||
secret_name=KeyManagementEventHooks._get_secret_name(
|
||||
key.key_alias
|
||||
)
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
|
@ -5,7 +5,9 @@ model_list:
|
|||
api_base: https://example-openai-endpoint.onrender.com/chat/completions
|
||||
api_key: "ishaan"
|
||||
|
||||
general_settings:
|
||||
master_key: sk-1234
|
||||
alerting: ["slack"]
|
||||
alerting_threshold: 0.0000001
|
||||
general_settings:
|
||||
key_management_system: "hashicorp_vault" # 👈 KEY CHANGE
|
||||
key_management_settings:
|
||||
store_virtual_keys: true # OPTIONAL. Defaults to False, when True will store virtual keys in secret manager
|
||||
prefix_for_stored_virtual_keys: "litellm/" # OPTIONAL. If set, this prefix will be used for stored virtual keys in the secret manager
|
||||
access_mode: "write_only" # Literal["read_only", "write_only", "read_and_write"]
|
|
@ -29,8 +29,10 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
from litellm.proxy._types import KeyManagementSystem
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
from .base_secret_manager import BaseSecretManager
|
||||
|
||||
class AWSSecretsManagerV2(BaseAWSLLM):
|
||||
|
||||
class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
||||
@classmethod
|
||||
def validate_environment(cls):
|
||||
if "AWS_REGION_NAME" not in os.environ:
|
||||
|
@ -146,7 +148,6 @@ class AWSSecretsManagerV2(BaseAWSLLM):
|
|||
secret_name: str,
|
||||
secret_value: str,
|
||||
description: Optional[str] = None,
|
||||
client_request_token: Optional[str] = None,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> dict:
|
||||
|
@ -157,7 +158,6 @@ class AWSSecretsManagerV2(BaseAWSLLM):
|
|||
secret_name: Name of the secret
|
||||
secret_value: Value to store (can be a JSON string)
|
||||
description: Optional description for the secret
|
||||
client_request_token: Optional unique identifier to ensure idempotency
|
||||
optional_params: Additional AWS parameters
|
||||
timeout: Request timeout
|
||||
"""
|
||||
|
|
95
litellm/secret_managers/base_secret_manager.py
Normal file
95
litellm/secret_managers/base_secret_manager.py
Normal file
|
@ -0,0 +1,95 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
class BaseSecretManager(ABC):
|
||||
"""
|
||||
Abstract base class for secret management implementations.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def async_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Asynchronously read a secret from the secret manager.
|
||||
|
||||
Args:
|
||||
secret_name (str): Name/path of the secret to read
|
||||
optional_params (Optional[dict]): Additional parameters specific to the secret manager
|
||||
timeout (Optional[Union[float, httpx.Timeout]]): Request timeout
|
||||
|
||||
Returns:
|
||||
Optional[str]: The secret value if found, None otherwise
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def sync_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Synchronously read a secret from the secret manager.
|
||||
|
||||
Args:
|
||||
secret_name (str): Name/path of the secret to read
|
||||
optional_params (Optional[dict]): Additional parameters specific to the secret manager
|
||||
timeout (Optional[Union[float, httpx.Timeout]]): Request timeout
|
||||
|
||||
Returns:
|
||||
Optional[str]: The secret value if found, None otherwise
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def async_write_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
secret_value: str,
|
||||
description: Optional[str] = None,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Asynchronously write a secret to the secret manager.
|
||||
|
||||
Args:
|
||||
secret_name (str): Name/path of the secret to write
|
||||
secret_value (str): Value to store
|
||||
description (Optional[str]): Description of the secret. Some secret managers allow storing a description with the secret.
|
||||
optional_params (Optional[dict]): Additional parameters specific to the secret manager
|
||||
timeout (Optional[Union[float, httpx.Timeout]]): Request timeout
|
||||
Returns:
|
||||
Dict[str, Any]: Response from the secret manager containing write operation details
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def async_delete_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
recovery_window_in_days: Optional[int] = 7,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Async function to delete a secret from the secret manager
|
||||
|
||||
Args:
|
||||
secret_name: Name of the secret to delete
|
||||
recovery_window_in_days: Number of days before permanent deletion (default: 7)
|
||||
optional_params: Additional parameters specific to the secret manager
|
||||
timeout: Request timeout
|
||||
|
||||
Returns:
|
||||
dict: Response from the secret manager containing deletion details
|
||||
"""
|
||||
pass
|
|
@ -1,5 +1,5 @@
|
|||
import os
|
||||
from typing import Optional
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
|
@ -13,8 +13,10 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.proxy._types import KeyManagementSystem
|
||||
|
||||
from .base_secret_manager import BaseSecretManager
|
||||
|
||||
class HashicorpSecretManager:
|
||||
|
||||
class HashicorpSecretManager(BaseSecretManager):
|
||||
def __init__(self):
|
||||
from litellm.proxy.proxy_server import CommonProxyErrors, premium_user
|
||||
|
||||
|
@ -110,7 +112,12 @@ class HashicorpSecretManager:
|
|||
return {"X-Vault-Token": self._auth_via_tls_cert()}
|
||||
return {"X-Vault-Token": self.vault_token}
|
||||
|
||||
async def async_read_secret(self, secret_name: str) -> Optional[str]:
|
||||
async def async_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Reads a secret from Vault KV v2 using an async HTTPX client.
|
||||
secret_name is just the path inside the KV mount (e.g., 'myapp/config').
|
||||
|
@ -140,7 +147,12 @@ class HashicorpSecretManager:
|
|||
verbose_logger.exception(f"Error reading secret from Hashicorp Vault: {e}")
|
||||
return None
|
||||
|
||||
def read_secret(self, secret_name: str) -> Optional[str]:
|
||||
def sync_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Reads a secret from Vault KV v2 using a sync HTTPX client.
|
||||
secret_name is just the path inside the KV mount (e.g., 'myapp/config').
|
||||
|
@ -166,6 +178,95 @@ class HashicorpSecretManager:
|
|||
verbose_logger.exception(f"Error reading secret from Hashicorp Vault: {e}")
|
||||
return None
|
||||
|
||||
async def async_write_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
secret_value: str,
|
||||
description: Optional[str] = None,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Writes a secret to Vault KV v2 using an async HTTPX client.
|
||||
|
||||
Args:
|
||||
secret_name: Path inside the KV mount (e.g., 'myapp/config')
|
||||
secret_value: Value to store
|
||||
description: Optional description for the secret
|
||||
optional_params: Additional parameters to include in the secret data
|
||||
timeout: Request timeout
|
||||
|
||||
Returns:
|
||||
dict: Response containing status and details of the operation
|
||||
"""
|
||||
async_client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.SecretManager,
|
||||
params={"timeout": timeout},
|
||||
)
|
||||
|
||||
try:
|
||||
url = self.get_url(secret_name)
|
||||
|
||||
# Prepare the secret data
|
||||
data = {"data": {"key": secret_value}}
|
||||
|
||||
if description:
|
||||
data["data"]["description"] = description
|
||||
|
||||
response = await async_client.post(
|
||||
url=url, headers=self._get_request_headers(), json=data
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error writing secret to Hashicorp Vault: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
async def async_delete_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
recovery_window_in_days: Optional[int] = 7,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Async function to delete a secret from Hashicorp Vault.
|
||||
In KV v2, this marks the latest version of the secret as deleted.
|
||||
|
||||
Args:
|
||||
secret_name: Name of the secret to delete
|
||||
recovery_window_in_days: Not used for Vault (Vault handles this internally)
|
||||
optional_params: Additional parameters specific to the secret manager
|
||||
timeout: Request timeout
|
||||
|
||||
Returns:
|
||||
dict: Response containing status and details of the operation
|
||||
"""
|
||||
async_client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.SecretManager,
|
||||
params={"timeout": timeout},
|
||||
)
|
||||
|
||||
try:
|
||||
# For KV v2 delete: /v1/<mount>/data/<path>
|
||||
url = self.get_url(secret_name)
|
||||
|
||||
response = await async_client.delete(
|
||||
url=url, headers=self._get_request_headers()
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
# Clear the cache for this secret
|
||||
self.cache.delete_cache(secret_name)
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"Secret {secret_name} deleted successfully",
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error deleting secret from Hashicorp Vault: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
def _get_secret_value_from_json_response(
|
||||
self, json_resp: Optional[dict]
|
||||
) -> Optional[str]:
|
||||
|
|
|
@ -291,7 +291,7 @@ def get_secret( # noqa: PLR0915
|
|||
raise e
|
||||
elif key_manager == KeyManagementSystem.HASHICORP_VAULT.value:
|
||||
try:
|
||||
secret = client.read_secret(secret_name)
|
||||
secret = client.sync_read_secret(secret_name=secret_name)
|
||||
if secret is None:
|
||||
raise ValueError(
|
||||
f"No secret found in Hashicorp Secret Manager for {secret_name}"
|
||||
|
|
|
@ -3561,7 +3561,7 @@ async def test_key_generate_with_secret_manager_call(prisma_client):
|
|||
# read from the secret manager
|
||||
|
||||
result = await aws_secret_manager_client.async_read_secret(
|
||||
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key_alias}"
|
||||
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{key_alias}"
|
||||
)
|
||||
|
||||
# Assert the correct key is stored in the secret manager
|
||||
|
@ -3582,7 +3582,7 @@ async def test_key_generate_with_secret_manager_call(prisma_client):
|
|||
# Assert the key is deleted from the secret manager
|
||||
|
||||
result = await aws_secret_manager_client.async_read_secret(
|
||||
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key_alias}"
|
||||
secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{key_alias}"
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
|
|
@ -13,6 +13,7 @@ sys.path.insert(
|
|||
from unittest.mock import patch, MagicMock
|
||||
import logging
|
||||
from litellm._logging import verbose_logger
|
||||
import uuid
|
||||
|
||||
verbose_logger.setLevel(logging.DEBUG)
|
||||
|
||||
|
@ -42,6 +43,25 @@ mock_vault_response = {
|
|||
"mount_type": "kv",
|
||||
}
|
||||
|
||||
# Update the mock_vault_response for write operations
|
||||
mock_write_response = {
|
||||
"request_id": "80fafb6a-e96a-4c5b-29fa-ff505ac72201",
|
||||
"lease_id": "",
|
||||
"renewable": False,
|
||||
"lease_duration": 0,
|
||||
"data": {
|
||||
"created_time": "2025-01-04T16:58:42.684673531Z",
|
||||
"custom_metadata": None,
|
||||
"deletion_time": "",
|
||||
"destroyed": False,
|
||||
"version": 1,
|
||||
},
|
||||
"wrap_info": None,
|
||||
"warnings": None,
|
||||
"auth": None,
|
||||
"mount_type": "kv",
|
||||
}
|
||||
|
||||
|
||||
def test_hashicorp_secret_manager_get_secret():
|
||||
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") as mock_get:
|
||||
|
@ -52,7 +72,7 @@ def test_hashicorp_secret_manager_get_secret():
|
|||
mock_get.return_value = mock_response
|
||||
|
||||
# Test the secret manager
|
||||
secret = hashicorp_secret_manager.read_secret("sample-secret-mock")
|
||||
secret = hashicorp_secret_manager.sync_read_secret("sample-secret-mock")
|
||||
assert secret == "value-mock"
|
||||
|
||||
# Verify the request was made with correct parameters
|
||||
|
@ -67,6 +87,84 @@ def test_hashicorp_secret_manager_get_secret():
|
|||
assert "X-Vault-Token" in mock_get.call_args.kwargs["headers"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hashicorp_secret_manager_write_secret():
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post"
|
||||
) as mock_post:
|
||||
# Configure the mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = (
|
||||
mock_write_response # Use the write-specific response
|
||||
)
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
# Test the secret manager
|
||||
secret_name = f"sample-secret-test-{uuid.uuid4()}"
|
||||
secret_value = f"value-mock-{uuid.uuid4()}"
|
||||
response = await hashicorp_secret_manager.async_write_secret(
|
||||
secret_name=secret_name,
|
||||
secret_value=secret_value,
|
||||
)
|
||||
|
||||
# Verify the response and that the request was made correctly
|
||||
assert (
|
||||
response == mock_write_response
|
||||
) # Compare against write-specific response
|
||||
mock_post.assert_called_once()
|
||||
print("CALL ARGS=", mock_post.call_args)
|
||||
print("call args[1]=", mock_post.call_args[1])
|
||||
|
||||
# Verify URL
|
||||
called_url = mock_post.call_args[1]["url"]
|
||||
assert secret_name in called_url
|
||||
assert (
|
||||
called_url
|
||||
== f"{hashicorp_secret_manager.vault_addr}/v1/admin/secret/data/{secret_name}"
|
||||
)
|
||||
|
||||
# Verify request body
|
||||
json_data = mock_post.call_args[1]["json"]
|
||||
assert "data" in json_data
|
||||
assert "key" in json_data["data"]
|
||||
assert json_data["data"]["key"] == secret_value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hashicorp_secret_manager_delete_secret():
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.delete"
|
||||
) as mock_delete:
|
||||
# Configure the mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_delete.return_value = mock_response
|
||||
|
||||
# Test the secret manager
|
||||
secret_name = f"sample-secret-test-{uuid.uuid4()}"
|
||||
response = await hashicorp_secret_manager.async_delete_secret(
|
||||
secret_name=secret_name
|
||||
)
|
||||
|
||||
# Verify the response
|
||||
assert response == {
|
||||
"status": "success",
|
||||
"message": f"Secret {secret_name} deleted successfully",
|
||||
}
|
||||
|
||||
# Verify the request was made correctly
|
||||
mock_delete.assert_called_once()
|
||||
|
||||
# Verify URL
|
||||
called_url = mock_delete.call_args[1]["url"]
|
||||
assert secret_name in called_url
|
||||
assert (
|
||||
called_url
|
||||
== f"{hashicorp_secret_manager.vault_addr}/v1/admin/secret/data/{secret_name}"
|
||||
)
|
||||
|
||||
|
||||
def test_hashicorp_secret_manager_tls_cert_auth():
|
||||
with patch("httpx.post") as mock_post:
|
||||
# Configure the mock response for TLS auth
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue