diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index a8325d3461..c07f30f864 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -150,15 +150,26 @@ class KeyManagementEventHooks: existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}" ) + new_secret_name = ( + response.key_alias + or data.key_alias + or f"virtual-key-{response.token_id}" + ) + verbose_proxy_logger.info( + "Updating secret in secret manager: secret_name=%s", + new_secret_name, + ) team_id = getattr(existing_key_row, "team_id", None) await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager( current_secret_name=initial_secret_name, - new_secret_name=response.key_alias - or data.key_alias - or f"virtual-key-{response.token_id}", + new_secret_name=new_secret_name, new_secret_value=response.key, team_id=team_id, ) + verbose_proxy_logger.info( + "Secret updated in secret manager: secret_name=%s", + new_secret_name, + ) except Exception as e: verbose_proxy_logger.warning( f"Failed to rotate virtual key in secret manager: {e}" diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2e71759072..152b09a86b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2770,6 +2770,7 @@ async def can_modify_verification_token( Rules: - Proxy admin can modify any key + - Internal jobs service account can modify any key (for auto-rotation) - For team keys: only team admin or key owner can modify - For personal keys: only key owner can modify @@ -2782,13 +2783,19 @@ async def can_modify_verification_token( Returns: True if user can modify the key, False otherwise """ + from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + is_team_key = _is_team_key(data=key_info) # 1. Proxy admin can modify any key if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return True - # 2. For team keys: only team admin or key owner can modify + # 2. Internal jobs service account can modify any key (for auto-rotation) + if user_api_key_dict.api_key == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME: + return True + + # 3. For team keys: only team admin or key owner can modify if is_team_key and key_info.team_id is not None: # Get team object to check if user is team admin team_table = await get_team_object( @@ -2818,7 +2825,7 @@ async def can_modify_verification_token( # Not team admin and doesn't own the key return False - # 3. For personal keys: only key owner can modify + # 4. For personal keys: only key owner can modify if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: return True @@ -3179,7 +3186,7 @@ def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: dependencies=[Depends(user_api_key_auth)], ) @management_endpoint_wrapper -async def regenerate_key_fn( +async def regenerate_key_fn( # noqa: PLR0915 key: Optional[str] = None, data: Optional[RegenerateKeyRequest] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -3330,6 +3337,10 @@ async def regenerate_key_fn( detail={"error": "You are not authorized to regenerate this key"}, ) + verbose_proxy_logger.info( + "Key regeneration requested: key_alias=%s", + getattr(_key_in_db, "key_alias", None), + ) verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) new_token = get_new_token(data=data) @@ -3380,6 +3391,10 @@ async def regenerate_key_fn( **updated_token_dict, ) + verbose_proxy_logger.info( + "Key regeneration completed: key_alias=%s", + getattr(_key_in_db, "key_alias", None), + ) asyncio.create_task( KeyManagementEventHooks.async_key_rotated_hook( data=data, diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 8edfc48336..c1b4d019dc 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -3,7 +3,8 @@ This is a file for the AWS Secret Manager Integration Handles Async Operations for: - Read Secret -- Write Secret +- Write Secret (CreateSecret) +- Update Secret (PutSecretValue) - for in-place rotation when alias is preserved - Delete Secret Relevant issue: https://github.com/BerriAI/litellm/issues/1883 @@ -42,11 +43,11 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): aws_profile_name: Optional[str] = None, aws_web_identity_token: Optional[str] = None, aws_sts_endpoint: Optional[str] = None, - **kwargs + **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 @@ -61,7 +62,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): # 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 + "AWS_REGION_NAME" not in os.environ and "AWS_REGION" not in os.environ and "AWS_DEFAULT_REGION" not in os.environ ): @@ -83,22 +84,36 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): return try: cls.validate_environment() - + # 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), + "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 @@ -246,13 +261,13 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): return primary_secret_kv_pairs.get(secret_name) 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, - tags: Optional[Union[dict, list]] = None + self, + secret_name: str, + secret_value: str, + description: Optional[str] = None, + optional_params: Optional[dict] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + tags: Optional[Union[dict, list]] = None, ) -> dict: """ Async function to write a secret to AWS Secrets Manager @@ -312,6 +327,94 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): except httpx.TimeoutException: raise ValueError("Timeout error occurred") + async def async_put_secret_value( + self, + secret_name: str, + secret_value: str, + optional_params: Optional[dict] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + ) -> dict: + """ + Async function to update an existing secret's value in AWS Secrets Manager. + + Uses PutSecretValue to update in place. Use this when rotating a secret + that keeps the same name (current_secret_name == new_secret_name). + + Args: + secret_name: Name of the existing secret to update + secret_value: New value to store + optional_params: Additional AWS parameters + timeout: Request timeout + + Returns: + dict: Response from AWS Secrets Manager containing update details + """ + from litellm._uuid import uuid + + data: Dict[str, Any] = { + "SecretId": secret_name, + "SecretString": secret_value, + "ClientRequestToken": str(uuid.uuid4()), + } + + endpoint_url, headers, body = self._prepare_request( + action="PutSecretValue", + secret_name=secret_name, + secret_value=secret_value, + optional_params=optional_params, + request_data=data, + ) + + async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.SecretManager, + params={"timeout": timeout}, + ) + + try: + response = await async_client.post( + url=endpoint_url, headers=headers, data=body.decode("utf-8") + ) + response.raise_for_status() + return response.json() + except httpx.HTTPStatusError as err: + raise ValueError(f"HTTP error occurred: {err.response.text}") + except httpx.TimeoutException: + raise ValueError("Timeout error occurred") + + 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: + """ + Rotate a secret. When current_secret_name == new_secret_name (in-place + update), uses PutSecretValue instead of create+delete to avoid + ResourceExistsException. + """ + if current_secret_name == new_secret_name: + # Same alias: update in place via PutSecretValue + verbose_logger.info( + "Secret rotated in-place (PutSecretValue): secret_name=%s", + current_secret_name, + ) + return await self.async_put_secret_value( + secret_name=current_secret_name, + secret_value=new_secret_value, + optional_params=optional_params, + timeout=timeout, + ) + # Different names: create new, delete old (base class logic) + return await super().async_rotate_secret( + current_secret_name=current_secret_name, + new_secret_name=new_secret_name, + new_secret_value=new_secret_value, + optional_params=optional_params, + timeout=timeout, + ) + async def async_delete_secret( self, secret_name: str, @@ -375,7 +478,7 @@ 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: @@ -388,11 +491,14 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): 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: + 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 ) @@ -431,12 +537,3 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): prepped = request.prepare() return endpoint_url, prepped.headers, body - - -# if __name__ == "__main__": -# print("loading aws secret manager v2") -# aws_secret_manager_v2 = AWSSecretsManagerV2() -# import asyncio -# print("writing secret to aws secret manager v2") -# asyncio.run(aws_secret_manager_v2.async_write_secret(secret_name="test_secret_3", secret_value="test_value_2")) -# print("reading secret from aws secret manager v2") diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py b/tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py new file mode 100644 index 0000000000..8398248262 --- /dev/null +++ b/tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py @@ -0,0 +1,109 @@ +""" +Regression tests for AWS Secrets Manager same-name in-place rotation fix. + +When current_secret_name == new_secret_name (e.g. key alias preserved during +rotation), AWS must use PutSecretValue to update in place instead of +create+delete, which would fail with ResourceExistsException. +""" +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 + + +@pytest.mark.asyncio +async def test_rotate_secret_same_name_uses_put_secret_value(): + """ + When current_secret_name == new_secret_name, async_rotate_secret should + call PutSecretValue (async_put_secret_value) instead of create+delete. + """ + secret_name = "litellm/tenant/litellm-metis-key" + new_value = "sk-new-rotated-key-value" + + with patch.object( + AWSSecretsManagerV2, + "async_put_secret_value", + new_callable=AsyncMock, + return_value={"ARN": "arn:aws:secretsmanager:us-east-1:123:secret:test"}, + ) as mock_put: + with patch.object( + AWSSecretsManagerV2, + "async_write_secret", + new_callable=AsyncMock, + ) as mock_write: + with patch.object( + AWSSecretsManagerV2, + "async_delete_secret", + new_callable=AsyncMock, + ) as mock_delete: + manager = AWSSecretsManagerV2() + result = await manager.async_rotate_secret( + current_secret_name=secret_name, + new_secret_name=secret_name, + new_secret_value=new_value, + ) + + # PutSecretValue (in-place update) should be called + mock_put.assert_called_once_with( + secret_name=secret_name, + secret_value=new_value, + optional_params=None, + timeout=None, + ) + # Create + delete should NOT be called + mock_write.assert_not_called() + mock_delete.assert_not_called() + assert result["ARN"] == "arn:aws:secretsmanager:us-east-1:123:secret:test" + + +@pytest.mark.asyncio +async def test_rotate_secret_different_names_uses_create_delete(): + """ + When current_secret_name != new_secret_name, async_rotate_secret should + use base class logic (create new, delete old). + """ + current_name = "litellm/old-key-alias" + new_name = "litellm/virtual-key-new-token-id" + new_value = "sk-new-key-value" + + with patch.object( + AWSSecretsManagerV2, + "async_read_secret", + new_callable=AsyncMock, + side_effect=["sk-old-value", new_value], # read old, then read new + ): + with patch.object( + AWSSecretsManagerV2, + "async_write_secret", + new_callable=AsyncMock, + return_value={"ARN": "arn:new"}, + ) as mock_write: + with patch.object( + AWSSecretsManagerV2, + "async_delete_secret", + new_callable=AsyncMock, + return_value={}, + ) as mock_delete: + with patch.object( + AWSSecretsManagerV2, + "async_put_secret_value", + new_callable=AsyncMock, + ) as mock_put: + manager = AWSSecretsManagerV2() + await manager.async_rotate_secret( + current_secret_name=current_name, + new_secret_name=new_name, + new_secret_value=new_value, + ) + + # PutSecretValue should NOT be called (different names) + mock_put.assert_not_called() + # Create + delete should be called + mock_write.assert_called_once() + mock_delete.assert_called_once_with( + secret_name=current_name, + recovery_window_in_days=7, + optional_params=None, + timeout=None, + )