diff --git a/docs/my-website/docs/secret_managers/aws_secret_manager.md b/docs/my-website/docs/secret_managers/aws_secret_manager.md index 44fa23a4ae..5b7ab1e3e7 100644 --- a/docs/my-website/docs/secret_managers/aws_secret_manager.md +++ b/docs/my-website/docs/secret_managers/aws_secret_manager.md @@ -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 | + + + diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4040b0aa70..9c98ad215b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 ( diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 8f3547fd08..8edfc48336 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -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 ) diff --git a/litellm/types/secret_managers/main.py b/litellm/types/secret_managers/main.py index 8d7aa764d6..e4c7d76573 100644 --- a/litellm/types/secret_managers/main.py +++ b/litellm/types/secret_managers/main.py @@ -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" - """ \ No newline at end of file + """ + + # 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)""" \ No newline at end of file diff --git a/tests/litellm_utils_tests/test_aws_secret_manager.py b/tests/litellm_utils_tests/test_aws_secret_manager.py index b6ccd41768..3870d336f0 100644 --- a/tests/litellm_utils_tests/test_aws_secret_manager.py +++ b/tests/litellm_utils_tests/test_aws_secret_manager.py @@ -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}")