[Feat] Adds IAM role assumption support for AWS Secret Manager (#16887)
* add AWS fields for KeyManagementSettings * docs IAM roles * use aws iam auth on secret manager v2 * fix: load_aws_secret_manager * test_secret_manager_with_iam_role_settings
This commit is contained in:
Ishaan Jaff
GitHub
parent
9120a02474
commit
57544f1662
@ -110,3 +110,57 @@ The `primary_secret_name` allows you to read multiple keys from a single AWS Sec
|
||||
|
||||
This reduces the number of AWS Secrets you need to manage.
|
||||
|
||||
## IAM Role Assumption
|
||||
|
||||
Use IAM roles instead of static AWS credentials for better security.
|
||||
|
||||
### Basic IAM Role
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
key_management_system: "aws_secret_manager"
|
||||
key_management_settings:
|
||||
store_virtual_keys: true
|
||||
aws_region_name: "us-east-1"
|
||||
aws_role_name: "arn:aws:iam::123456789012:role/LiteLLMSecretManagerRole"
|
||||
aws_session_name: "litellm-session"
|
||||
```
|
||||
|
||||
### Cross-Account Access
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
key_management_system: "aws_secret_manager"
|
||||
key_management_settings:
|
||||
store_virtual_keys: true
|
||||
aws_region_name: "us-east-1"
|
||||
aws_role_name: "arn:aws:iam::999999999999:role/CrossAccountRole"
|
||||
aws_external_id: "unique-external-id"
|
||||
```
|
||||
|
||||
### EKS with IRSA
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
key_management_system: "aws_secret_manager"
|
||||
key_management_settings:
|
||||
store_virtual_keys: true
|
||||
aws_region_name: "us-east-1"
|
||||
aws_role_name: "arn:aws:iam::123456789012:role/LiteLLMServiceAccountRole"
|
||||
aws_web_identity_token: "os.environ/AWS_WEB_IDENTITY_TOKEN_FILE"
|
||||
```
|
||||
|
||||
### Configuration Parameters
|
||||
|
||||
| Parameter | Description |
|
||||
|-----------|-------------|
|
||||
| `aws_region_name` | AWS region |
|
||||
| `aws_role_name` | IAM role ARN to assume |
|
||||
| `aws_session_name` | Session name (optional) |
|
||||
| `aws_external_id` | External ID for cross-account |
|
||||
| `aws_profile_name` | AWS profile from `~/.aws/credentials` |
|
||||
| `aws_web_identity_token` | OIDC token path for IRSA |
|
||||
| `aws_sts_endpoint` | Custom STS endpoint for VPC |
|
||||
|
||||
|
||||
|
||||
|
||||
@ -2692,7 +2692,10 @@ class ProxyConfig:
|
||||
AWSSecretsManagerV2,
|
||||
)
|
||||
|
||||
AWSSecretsManagerV2.load_aws_secret_manager(use_aws_secret_manager=True)
|
||||
AWSSecretsManagerV2.load_aws_secret_manager(
|
||||
use_aws_secret_manager=True,
|
||||
key_management_settings=litellm._key_management_settings,
|
||||
)
|
||||
elif key_management_system == KeyManagementSystem.AWS_KMS.value:
|
||||
load_aws_kms(use_aws_kms=True)
|
||||
elif (
|
||||
|
||||
@ -33,25 +33,73 @@ from .base_secret_manager import BaseSecretManager
|
||||
|
||||
|
||||
class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
||||
def __init__(self, **kwargs):
|
||||
def __init__(
|
||||
self,
|
||||
aws_region_name: Optional[str] = None,
|
||||
aws_role_name: Optional[str] = None,
|
||||
aws_session_name: Optional[str] = None,
|
||||
aws_external_id: Optional[str] = None,
|
||||
aws_profile_name: Optional[str] = None,
|
||||
aws_web_identity_token: Optional[str] = None,
|
||||
aws_sts_endpoint: Optional[str] = None,
|
||||
**kwargs
|
||||
):
|
||||
BaseSecretManager.__init__(self, **kwargs)
|
||||
BaseAWSLLM.__init__(self, **kwargs)
|
||||
|
||||
# Store AWS authentication settings
|
||||
self.aws_region_name = aws_region_name
|
||||
self.aws_role_name = aws_role_name
|
||||
self.aws_session_name = aws_session_name
|
||||
self.aws_external_id = aws_external_id
|
||||
self.aws_profile_name = aws_profile_name
|
||||
self.aws_web_identity_token = aws_web_identity_token
|
||||
self.aws_sts_endpoint = aws_sts_endpoint
|
||||
|
||||
@classmethod
|
||||
def validate_environment(cls):
|
||||
if "AWS_REGION_NAME" not in os.environ:
|
||||
raise ValueError("Missing required environment variable - AWS_REGION_NAME")
|
||||
# AWS_REGION_NAME is only strictly required if not using a profile or role
|
||||
# When using IAM roles, the region can come from multiple sources
|
||||
if (
|
||||
"AWS_REGION_NAME" not in os.environ
|
||||
and "AWS_REGION" not in os.environ
|
||||
and "AWS_DEFAULT_REGION" not in os.environ
|
||||
):
|
||||
verbose_logger.warning(
|
||||
"No AWS region found in environment. Ensure aws_region_name is set in key_management_settings "
|
||||
"or AWS_REGION_NAME/AWS_REGION/AWS_DEFAULT_REGION is set in environment."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def load_aws_secret_manager(cls, use_aws_secret_manager: Optional[bool]):
|
||||
def load_aws_secret_manager(
|
||||
cls,
|
||||
use_aws_secret_manager: Optional[bool],
|
||||
key_management_settings: Optional[Any] = None,
|
||||
):
|
||||
"""
|
||||
Initialize AWSSecretsManagerV2 and sets litellm.secret_manager_client = AWSSecretsManagerV2() and litellm._key_management_system = KeyManagementSystem.AWS_SECRET_MANAGER
|
||||
Initialize AWSSecretsManagerV2 with settings from key_management_settings
|
||||
"""
|
||||
if use_aws_secret_manager is None or use_aws_secret_manager is False:
|
||||
return
|
||||
try:
|
||||
cls.validate_environment()
|
||||
litellm.secret_manager_client = cls()
|
||||
|
||||
# Extract AWS settings from key_management_settings if provided
|
||||
aws_kwargs = {}
|
||||
if key_management_settings is not None:
|
||||
aws_kwargs = {
|
||||
"aws_region_name": getattr(key_management_settings, "aws_region_name", None),
|
||||
"aws_role_name": getattr(key_management_settings, "aws_role_name", None),
|
||||
"aws_session_name": getattr(key_management_settings, "aws_session_name", None),
|
||||
"aws_external_id": getattr(key_management_settings, "aws_external_id", None),
|
||||
"aws_profile_name": getattr(key_management_settings, "aws_profile_name", None),
|
||||
"aws_web_identity_token": getattr(key_management_settings, "aws_web_identity_token", None),
|
||||
"aws_sts_endpoint": getattr(key_management_settings, "aws_sts_endpoint", None),
|
||||
}
|
||||
# Remove None values
|
||||
aws_kwargs = {k: v for k, v in aws_kwargs.items() if v is not None}
|
||||
|
||||
litellm.secret_manager_client = cls(**aws_kwargs)
|
||||
litellm._key_management_system = KeyManagementSystem.AWS_SECRET_MANAGER
|
||||
|
||||
except Exception as e:
|
||||
@ -327,6 +375,24 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
optional_params = optional_params or {}
|
||||
|
||||
# Build optional_params from instance settings if not provided
|
||||
# This allows the IAM role settings to be used for Secret Manager calls
|
||||
if not optional_params.get("aws_role_name") and self.aws_role_name:
|
||||
optional_params["aws_role_name"] = self.aws_role_name
|
||||
if not optional_params.get("aws_session_name") and self.aws_session_name:
|
||||
optional_params["aws_session_name"] = self.aws_session_name
|
||||
if not optional_params.get("aws_region_name") and self.aws_region_name:
|
||||
optional_params["aws_region_name"] = self.aws_region_name
|
||||
if not optional_params.get("aws_external_id") and self.aws_external_id:
|
||||
optional_params["aws_external_id"] = self.aws_external_id
|
||||
if not optional_params.get("aws_profile_name") and self.aws_profile_name:
|
||||
optional_params["aws_profile_name"] = self.aws_profile_name
|
||||
if not optional_params.get("aws_web_identity_token") and self.aws_web_identity_token:
|
||||
optional_params["aws_web_identity_token"] = self.aws_web_identity_token
|
||||
if not optional_params.get("aws_sts_endpoint") and self.aws_sts_endpoint:
|
||||
optional_params["aws_sts_endpoint"] = self.aws_sts_endpoint
|
||||
|
||||
boto3_credentials_info = self._get_boto_credentials_from_optional_params(
|
||||
optional_params
|
||||
)
|
||||
|
||||
@ -49,4 +49,26 @@ class KeyManagementSettings(LiteLLMPydanticObjectBase):
|
||||
"""
|
||||
Path to custom secret manager class (e.g. "my_secret_manager.InMemorySecretManager")
|
||||
Required when key_management_system is "custom"
|
||||
"""
|
||||
"""
|
||||
|
||||
# AWS IAM Role Assumption Settings (for AWS Secret Manager)
|
||||
aws_region_name: Optional[str] = None
|
||||
"""AWS region for Secret Manager operations (e.g., 'us-east-1')"""
|
||||
|
||||
aws_role_name: Optional[str] = None
|
||||
"""ARN of IAM role to assume for Secret Manager access (e.g., 'arn:aws:iam::123456789012:role/MyRole')"""
|
||||
|
||||
aws_session_name: Optional[str] = None
|
||||
"""Session name for the assumed role session (optional, auto-generated if not provided)"""
|
||||
|
||||
aws_external_id: Optional[str] = None
|
||||
"""External ID for role assumption (required for cross-account access)"""
|
||||
|
||||
aws_profile_name: Optional[str] = None
|
||||
"""AWS profile name to use from ~/.aws/credentials"""
|
||||
|
||||
aws_web_identity_token: Optional[str] = None
|
||||
"""Web identity token for OIDC/IRSA authentication"""
|
||||
|
||||
aws_sts_endpoint: Optional[str] = None
|
||||
"""Custom STS endpoint URL (useful for VPC endpoints or testing)"""
|
||||
@ -31,6 +31,7 @@ import pytest
|
||||
from litellm._uuid import uuid
|
||||
import json
|
||||
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings
|
||||
|
||||
|
||||
def check_aws_credentials():
|
||||
@ -253,3 +254,211 @@ async def test_write_secret_with_description_and_tags():
|
||||
delete_response = await secret_manager.async_delete_secret(secret_name=test_secret_name)
|
||||
print("Delete Response:", delete_response)
|
||||
assert delete_response is not None
|
||||
|
||||
|
||||
def test_secret_manager_with_iam_role_settings():
|
||||
"""
|
||||
Test AWS Secret Manager initialization with IAM role settings
|
||||
"""
|
||||
settings = KeyManagementSettings(
|
||||
aws_region_name="us-east-1",
|
||||
aws_role_name="arn:aws:iam::123456789012:role/TestRole",
|
||||
aws_session_name="test-session",
|
||||
)
|
||||
|
||||
secret_manager = AWSSecretsManagerV2(
|
||||
aws_region_name=settings.aws_region_name,
|
||||
aws_role_name=settings.aws_role_name,
|
||||
aws_session_name=settings.aws_session_name,
|
||||
)
|
||||
|
||||
# Verify settings are stored
|
||||
assert secret_manager.aws_role_name == settings.aws_role_name
|
||||
assert secret_manager.aws_region_name == settings.aws_region_name
|
||||
assert secret_manager.aws_session_name == settings.aws_session_name
|
||||
|
||||
|
||||
def test_secret_manager_with_cross_account_settings():
|
||||
"""
|
||||
Test AWS Secret Manager initialization with cross-account IAM role settings
|
||||
"""
|
||||
settings = KeyManagementSettings(
|
||||
aws_region_name="us-west-2",
|
||||
aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole",
|
||||
aws_session_name="cross-account-session",
|
||||
aws_external_id="unique-external-id",
|
||||
)
|
||||
|
||||
secret_manager = AWSSecretsManagerV2(
|
||||
aws_region_name=settings.aws_region_name,
|
||||
aws_role_name=settings.aws_role_name,
|
||||
aws_session_name=settings.aws_session_name,
|
||||
aws_external_id=settings.aws_external_id,
|
||||
)
|
||||
|
||||
# Verify settings are stored
|
||||
assert secret_manager.aws_role_name == settings.aws_role_name
|
||||
assert secret_manager.aws_region_name == settings.aws_region_name
|
||||
assert secret_manager.aws_external_id == settings.aws_external_id
|
||||
|
||||
|
||||
def test_secret_manager_with_irsa_settings():
|
||||
"""
|
||||
Test AWS Secret Manager initialization with IRSA (EKS) settings
|
||||
"""
|
||||
settings = KeyManagementSettings(
|
||||
aws_region_name="us-east-1",
|
||||
aws_role_name="arn:aws:iam::123456789012:role/EKSServiceAccountRole",
|
||||
aws_session_name="eks-session",
|
||||
aws_web_identity_token="os.environ/AWS_WEB_IDENTITY_TOKEN_FILE",
|
||||
)
|
||||
|
||||
secret_manager = AWSSecretsManagerV2(
|
||||
aws_region_name=settings.aws_region_name,
|
||||
aws_role_name=settings.aws_role_name,
|
||||
aws_session_name=settings.aws_session_name,
|
||||
aws_web_identity_token=settings.aws_web_identity_token,
|
||||
)
|
||||
|
||||
# Verify settings are stored
|
||||
assert secret_manager.aws_role_name == settings.aws_role_name
|
||||
assert secret_manager.aws_web_identity_token == settings.aws_web_identity_token
|
||||
|
||||
|
||||
def test_secret_manager_with_custom_sts_endpoint():
|
||||
"""
|
||||
Test AWS Secret Manager initialization with custom STS endpoint (VPC endpoint)
|
||||
"""
|
||||
settings = KeyManagementSettings(
|
||||
aws_region_name="us-east-1",
|
||||
aws_role_name="arn:aws:iam::123456789012:role/VPCRole",
|
||||
aws_session_name="vpc-session",
|
||||
aws_sts_endpoint="https://sts.us-east-1.vpce-0123456789abcdef.amazonaws.com",
|
||||
)
|
||||
|
||||
secret_manager = AWSSecretsManagerV2(
|
||||
aws_region_name=settings.aws_region_name,
|
||||
aws_role_name=settings.aws_role_name,
|
||||
aws_session_name=settings.aws_session_name,
|
||||
aws_sts_endpoint=settings.aws_sts_endpoint,
|
||||
)
|
||||
|
||||
# Verify settings are stored
|
||||
assert secret_manager.aws_role_name == settings.aws_role_name
|
||||
assert secret_manager.aws_sts_endpoint == settings.aws_sts_endpoint
|
||||
|
||||
|
||||
def test_secret_manager_with_aws_profile():
|
||||
"""
|
||||
Test AWS Secret Manager initialization with AWS profile
|
||||
"""
|
||||
settings = KeyManagementSettings(
|
||||
aws_region_name="us-east-1",
|
||||
aws_profile_name="litellm-dev",
|
||||
)
|
||||
|
||||
secret_manager = AWSSecretsManagerV2(
|
||||
aws_region_name=settings.aws_region_name,
|
||||
aws_profile_name=settings.aws_profile_name,
|
||||
)
|
||||
|
||||
# Verify settings are stored
|
||||
assert secret_manager.aws_profile_name == settings.aws_profile_name
|
||||
|
||||
|
||||
def test_load_aws_secret_manager_with_settings():
|
||||
"""
|
||||
Test loading AWS Secret Manager with key_management_settings
|
||||
"""
|
||||
import litellm
|
||||
|
||||
settings = KeyManagementSettings(
|
||||
store_virtual_keys=True,
|
||||
aws_region_name="us-east-1",
|
||||
aws_role_name="arn:aws:iam::123456789012:role/TestRole",
|
||||
aws_session_name="test-session",
|
||||
)
|
||||
|
||||
# Set environment variable for validation to pass
|
||||
os.environ["AWS_REGION_NAME"] = "us-east-1"
|
||||
|
||||
try:
|
||||
AWSSecretsManagerV2.load_aws_secret_manager(
|
||||
use_aws_secret_manager=True,
|
||||
key_management_settings=settings,
|
||||
)
|
||||
|
||||
# Verify the client was created
|
||||
assert litellm.secret_manager_client is not None
|
||||
assert isinstance(litellm.secret_manager_client, AWSSecretsManagerV2)
|
||||
|
||||
# Verify settings were passed through
|
||||
assert litellm.secret_manager_client.aws_role_name == settings.aws_role_name
|
||||
assert litellm.secret_manager_client.aws_region_name == settings.aws_region_name
|
||||
assert litellm.secret_manager_client.aws_session_name == settings.aws_session_name
|
||||
finally:
|
||||
# Cleanup
|
||||
litellm.secret_manager_client = None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_to_end_iam_role_secret_write():
|
||||
"""
|
||||
Test writing a secret using IAM role assumption (integration test)
|
||||
|
||||
Requires:
|
||||
- AWS_REGION_NAME environment variable
|
||||
- TEST_IAM_ROLE_ARN environment variable with ARN of a role that can be assumed
|
||||
- Proper AWS credentials configured (via instance profile, IAM role, or environment)
|
||||
"""
|
||||
# Skip if TEST_IAM_ROLE_ARN is not set
|
||||
test_role_arn = os.getenv("TEST_IAM_ROLE_ARN")
|
||||
if not test_role_arn:
|
||||
pytest.skip("TEST_IAM_ROLE_ARN environment variable not set")
|
||||
|
||||
aws_region = os.getenv("AWS_REGION_NAME", "us-east-1")
|
||||
|
||||
settings = KeyManagementSettings(
|
||||
store_virtual_keys=True,
|
||||
aws_region_name=aws_region,
|
||||
aws_role_name=test_role_arn,
|
||||
aws_session_name="integration-test-session",
|
||||
)
|
||||
|
||||
secret_manager = AWSSecretsManagerV2(
|
||||
aws_region_name=settings.aws_region_name,
|
||||
aws_role_name=settings.aws_role_name,
|
||||
aws_session_name=settings.aws_session_name,
|
||||
)
|
||||
|
||||
test_secret_name = f"litellm_test_iam_{uuid.uuid4().hex[:8]}"
|
||||
test_secret_value = "test_value_iam_role"
|
||||
|
||||
try:
|
||||
# Test write operation using IAM role
|
||||
response = await secret_manager.async_write_secret(
|
||||
secret_name=test_secret_name,
|
||||
secret_value=test_secret_value,
|
||||
)
|
||||
|
||||
print("Write Response with IAM Role:", response)
|
||||
assert response is not None
|
||||
assert "ARN" in response
|
||||
|
||||
# Test read operation using IAM role
|
||||
read_value = await secret_manager.async_read_secret(
|
||||
secret_name=test_secret_name
|
||||
)
|
||||
|
||||
print("Read Value with IAM Role:", read_value)
|
||||
assert read_value == test_secret_value
|
||||
|
||||
finally:
|
||||
# Cleanup: Delete the secret
|
||||
try:
|
||||
delete_response = await secret_manager.async_delete_secret(
|
||||
secret_name=test_secret_name
|
||||
)
|
||||
print("Delete Response:", delete_response)
|
||||
except Exception as e:
|
||||
print(f"Cleanup failed: {e}")
|
||||
|
||||
Loading…
Reference in New Issue
Block a user