diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 0fb59105b5f..3e9f4e259d5 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -16,6 +16,8 @@ Requires: import json import os +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx @@ -294,32 +296,13 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): raise ValueError("Tags must be a dict or list of {Key, Value} pairs") data["Tags"] = tags_list - endpoint_url, headers, body = self._prepare_request( - action="CreateSecret", + create_response: Final = await self._async_create_or_restore_secret( secret_name=secret_name, - secret_value=secret_value, - optional_params=optional_params, request_data=data, + optional_params=optional_params, + timeout=timeout, ) - async_client: Final = get_async_httpx_client( - llm_provider=httpxSpecialProvider.SecretManager, - params={"timeout": timeout}, - ) - - try: - response: Final = await async_client.post( - url=endpoint_url, - headers=headers, - data=body.decode("utf-8"), - ) - response.raise_for_status() - create_response: Final = response.json() - except httpx.HTTPStatusError as err: - raise ValueError(f"HTTP error occurred: {err.response.text}") - except httpx.TimeoutException: - raise ValueError("Timeout error occurred") - if self.replica_regions: try: await self.async_replicate_secret( @@ -343,6 +326,110 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): return create_response + async def _async_create_or_restore_secret( + self, + secret_name: str, + request_data: Mapping[str, object], + optional_params: dict | None, + timeout: float | httpx.Timeout | None, + ) -> dict[str, object]: + try: + return await self._async_post_action( + action="CreateSecret", + secret_name=secret_name, + request_data=request_data, + optional_params=optional_params, + timeout=timeout, + ) + except ValueError: + if not await self._async_is_scheduled_for_deletion( + secret_name=secret_name, + optional_params=optional_params, + timeout=timeout, + ): + raise + + verbose_logger.info( + "Secret %s is scheduled for deletion, restoring and updating in place (RestoreSecret + UpdateSecret)", + secret_name, + ) + await self._async_post_action( + action="RestoreSecret", + secret_name=secret_name, + request_data=None, + optional_params=optional_params, + timeout=timeout, + ) + update_data: Final = MappingProxyType( + {("SecretId" if key == "Name" else key): value for key, value in request_data.items() if key != "Tags"} + ) + tags: Final = request_data.get("Tags") + try: + updated: Final = await self._async_post_action( + action="UpdateSecret", + secret_name=secret_name, + request_data=update_data, + optional_params=optional_params, + timeout=timeout, + ) + if tags is not None: + await self._async_post_action( + action="TagResource", + secret_name=secret_name, + request_data=MappingProxyType({"SecretId": secret_name, "Tags": tags}), + optional_params=optional_params, + timeout=timeout, + ) + except ValueError: + await self.async_delete_secret(secret_name=secret_name, optional_params=optional_params, timeout=timeout) + raise + return updated + + async def _async_is_scheduled_for_deletion( + self, + secret_name: str, + optional_params: dict | None, + timeout: float | httpx.Timeout | None, + ) -> bool: + try: + described: Final = await self._async_post_action( + action="DescribeSecret", + secret_name=secret_name, + request_data=None, + optional_params=optional_params, + timeout=timeout, + ) + except ValueError: + return False + return described.get("DeletedDate") is not None + + async def _async_post_action( + self, + action: str, + secret_name: str, + request_data: Mapping[str, object] | None, + optional_params: dict | None, + timeout: float | httpx.Timeout | None, + ) -> dict[str, object]: + endpoint_url, headers, body = self._prepare_request( + action=action, + secret_name=secret_name, + optional_params=optional_params, + request_data=dict(request_data) if request_data is not None else None, + ) + async_client: Final = get_async_httpx_client( + llm_provider=httpxSpecialProvider.SecretManager, + params={"timeout": timeout}, + ) + try: + response: Final = 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_replicate_secret( self, secret_name: str, 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 index 4a4cec6bf77..af1380f6fd1 100644 --- a/tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py +++ b/tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py @@ -1,11 +1,17 @@ -from collections.abc import Mapping +import json +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from dataclasses import dataclass, replace from types import MappingProxyType from typing import Final, TypeAlias +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 +from litellm.types.llms.custom_http import httpxSpecialProvider OptionalParams: TypeAlias = Mapping[str, object] | None @@ -204,3 +210,173 @@ async def test_rotate_secret_different_names_persists_requested_value_and_delete assert manager.storage.values[new_name] == new_value assert current_name not in manager.storage.values assert manager.storage.values[unrelated_secret_name] == unrelated_value + + +@dataclass(frozen=True, slots=True) +class FakeSecretsManagerState: + live: Mapping[str, str] + scheduled_for_deletion: frozenset[str] = frozenset() + actions: tuple[str, ...] = () + descriptions: Mapping[str, str] = MappingProxyType({}) + failing_actions: frozenset[str] = frozenset() + + +class FakeSecretsManagerService: + def __init__(self, state: FakeSecretsManagerState) -> None: + self.state = state + + def handle(self, request: httpx.Request) -> httpx.Response: + action: Final = request.headers["X-Amz-Target"].removeprefix("secretsmanager.") + body: Final = json.loads(request.content) + name: Final = str(body.get("Name") or body.get("SecretId")) + self.state = replace(self.state, actions=(*self.state.actions, f"{action}:{name}")) + if action in self.state.failing_actions: + return self._error("InternalServiceError", f"injected failure for {action}") + match action: + case "CreateSecret": + if name in self.state.live: + return self._error("ResourceExistsException", f"The secret {name} already exists") + self.state = replace( + self.state, + live=MappingProxyType({**self.state.live, name: str(body["SecretString"])}), + descriptions=MappingProxyType({**self.state.descriptions, name: str(body.get("Description", ""))}), + ) + return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name}) + case "UpdateSecret": + if name in self.state.scheduled_for_deletion: + return self._error( + "InvalidRequestException", + "You can't perform this operation on the secret because it was marked for deletion.", + ) + self.state = replace( + self.state, + live=MappingProxyType({**self.state.live, name: str(body["SecretString"])}), + descriptions=MappingProxyType({**self.state.descriptions, name: str(body.get("Description", ""))}), + ) + return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name}) + case "DescribeSecret": + if name not in self.state.live: + return self._error("ResourceNotFoundException", "Secrets Manager can't find the specified secret.") + deleted: Final = "2026-01-01T00:00:00Z" if name in self.state.scheduled_for_deletion else None + return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name, "DeletedDate": deleted}) + case "RestoreSecret": + self.state = replace(self.state, scheduled_for_deletion=self.state.scheduled_for_deletion - {name}) + return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name}) + case "PutSecretValue": + if name in self.state.scheduled_for_deletion: + return self._error( + "InvalidRequestException", + "You can't perform this operation on the secret because it was marked for deletion.", + ) + self.state = replace( + self.state, live=MappingProxyType({**self.state.live, name: str(body["SecretString"])}) + ) + return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name}) + case "GetSecretValue": + if name not in self.state.live or name in self.state.scheduled_for_deletion: + return self._error("ResourceNotFoundException", "Secrets Manager can't find the specified secret.") + return httpx.Response(200, json={"SecretString": self.state.live[name]}) + case "DeleteSecret": + self.state = replace(self.state, scheduled_for_deletion=self.state.scheduled_for_deletion | {name}) + return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name}) + return self._error("UnsupportedAction", action) + + @staticmethod + def _error(error_type: str, message: str) -> httpx.Response: + return httpx.Response(400, json={"__type": error_type, "message": message}) + + +@contextmanager +def fake_secrets_manager(monkeypatch: pytest.MonkeyPatch) -> Iterator[FakeSecretsManagerService]: + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "synthetic-access-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "synthetic-secret-key") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + service: Final = FakeSecretsManagerService(FakeSecretsManagerState(live=MappingProxyType({}))) + cache_key: Final = "async_httpx_clienttimeout_None" + httpxSpecialProvider.SecretManager + litellm.in_memory_llm_clients_cache.set_cache( + key=cache_key, + value=AsyncHTTPHandler(transport=httpx.MockTransport(service.handle)), + ) + try: + yield service + finally: + litellm.in_memory_llm_clients_cache.delete_cache( + litellm.in_memory_llm_clients_cache.update_cache_key_with_event_loop(cache_key) + ) + + +@pytest.mark.asyncio +async def test_rotate_secret_back_to_name_inside_recovery_window_restores_and_stores_new_value( + monkeypatch: pytest.MonkeyPatch, +) -> None: + alias_a: Final = "synthetic/alias-a" + alias_b: Final = "synthetic/alias-b" + with fake_secrets_manager(monkeypatch) as fake: + manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1") + await manager.async_write_secret(secret_name=alias_a, secret_value="value-1") + await manager.async_rotate_secret( + current_secret_name=alias_a, new_secret_name=alias_b, new_secret_value="value-2" + ) + assert alias_a in fake.state.scheduled_for_deletion + + await manager.async_rotate_secret( + current_secret_name=alias_b, new_secret_name=alias_a, new_secret_value="value-3" + ) + + assert await manager.async_read_secret(secret_name=alias_a) == "value-3" + assert await manager.async_read_secret(secret_name=alias_b) is None + assert fake.state.scheduled_for_deletion == frozenset({alias_b}) + assert fake.state.descriptions[alias_a] == f"Rotated from {alias_b}" + + +@pytest.mark.asyncio +async def test_write_secret_to_name_inside_recovery_window_reschedules_deletion_when_update_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + alias: Final = "synthetic/deleted-alias" + with fake_secrets_manager(monkeypatch) as fake: + manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1") + await manager.async_write_secret(secret_name=alias, secret_value="value-1") + await manager.async_delete_secret(secret_name=alias, recovery_window_in_days=7) + fake.state = replace(fake.state, failing_actions=frozenset({"UpdateSecret"})) + + with pytest.raises(ValueError, match="injected failure for UpdateSecret"): + await manager.async_write_secret(secret_name=alias, secret_value="value-2") + + assert fake.state.scheduled_for_deletion == frozenset({alias}) + assert fake.state.live[alias] == "value-1" + + +@pytest.mark.asyncio +async def test_write_secret_to_name_inside_recovery_window_restores_and_stores_new_value( + monkeypatch: pytest.MonkeyPatch, +) -> None: + alias: Final = "synthetic/deleted-alias" + with fake_secrets_manager(monkeypatch) as fake: + manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1") + await manager.async_write_secret(secret_name=alias, secret_value="value-1") + await manager.async_delete_secret(secret_name=alias, recovery_window_in_days=7) + + assert await manager.async_write_secret(secret_name=alias, secret_value="value-2") == { + "ARN": f"arn:fake:{alias}", + "Name": alias, + } + + assert await manager.async_read_secret(secret_name=alias) == "value-2" + assert fake.state.scheduled_for_deletion == frozenset() + + +@pytest.mark.asyncio +async def test_write_secret_to_live_existing_name_still_fails_without_overwriting( + monkeypatch: pytest.MonkeyPatch, +) -> None: + alias: Final = "synthetic/live-alias" + with fake_secrets_manager(monkeypatch) as fake: + manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1") + await manager.async_write_secret(secret_name=alias, secret_value="value-1") + + with pytest.raises(ValueError, match="ResourceExistsException"): + await manager.async_write_secret(secret_name=alias, secret_value="value-2") + + assert await manager.async_read_secret(secret_name=alias) == "value-1" + assert f"RestoreSecret:{alias}" not in fake.state.actions