From b36e704e06298184e44b720e489c7f93a1c5692b Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Wed, 21 Jan 2026 08:00:36 +0530 Subject: [PATCH] fix: ensure auto-rotation updates existing AWS secret instead of creating new one (#19455) --- .../common_utils/key_rotation_manager.py | 101 +++++++----- .../proxy/hooks/key_management_event_hooks.py | 3 +- .../test_key_rotation_integration.py | 144 ++++++++++++++++++ 3 files changed, 207 insertions(+), 41 deletions(-) create mode 100644 tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 3c367eafbc1..13bbf2272f7 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -26,97 +26,119 @@ class KeyRotationManager: """ Manages automated key rotation based on individual key rotation schedules. """ - + def __init__(self, prisma_client: PrismaClient): self.prisma_client = prisma_client - + async def process_rotations(self): """ Main entry point - find and rotate keys that are due for rotation """ try: verbose_proxy_logger.info("Starting scheduled key rotation check...") - + # Find keys that are due for rotation keys_to_rotate = await self._find_keys_needing_rotation() - + if not keys_to_rotate: verbose_proxy_logger.debug("No keys are due for rotation at this time") return - - verbose_proxy_logger.info(f"Found {len(keys_to_rotate)} keys due for rotation") - + + verbose_proxy_logger.info( + f"Found {len(keys_to_rotate)} keys due for rotation" + ) + # Rotate each key for key in keys_to_rotate: try: await self._rotate_key(key) - key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown") - verbose_proxy_logger.info(f"Successfully rotated key: {key_identifier}") + key_identifier = key.key_name or ( + key.token[:8] + "..." if key.token else "unknown" + ) + verbose_proxy_logger.info( + f"Successfully rotated key: {key_identifier}" + ) except Exception as e: - key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown") - verbose_proxy_logger.error(f"Failed to rotate key {key_identifier}: {e}") - + key_identifier = key.key_name or ( + key.token[:8] + "..." if key.token else "unknown" + ) + verbose_proxy_logger.error( + f"Failed to rotate key {key_identifier}: {e}" + ) + except Exception as e: verbose_proxy_logger.error(f"Key rotation process failed: {e}") - + async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]: """ Find keys that are due for rotation based on their key_rotation_at timestamp. - + Logic: - Key has auto_rotate = true - key_rotation_at is null (needs initial setup) OR key_rotation_at <= now """ now = datetime.now(timezone.utc) - - keys_with_rotation = await self.prisma_client.db.litellm_verificationtoken.find_many( - where={ - "auto_rotate": True, # Only keys marked for auto rotation - "OR": [ - {"key_rotation_at": None}, # Keys that need initial rotation time setup - {"key_rotation_at": {"lte": now}} # Keys where rotation time has passed - ] - } + + keys_with_rotation = ( + await self.prisma_client.db.litellm_verificationtoken.find_many( + where={ + "auto_rotate": True, # Only keys marked for auto rotation + "OR": [ + { + "key_rotation_at": None + }, # Keys that need initial rotation time setup + { + "key_rotation_at": {"lte": now} + }, # Keys where rotation time has passed + ], + } + ) ) - + return keys_with_rotation - + def _should_rotate_key(self, key: LiteLLM_VerificationToken, now: datetime) -> bool: """ Determine if a key should be rotated based on key_rotation_at timestamp. """ if not key.rotation_interval: return False - + # If key_rotation_at is not set, rotate immediately (and set it) if key.key_rotation_at is None: return True - + # Check if the rotation time has passed return now >= key.key_rotation_at - + async def _rotate_key(self, key: LiteLLM_VerificationToken): """ Rotate a single key using existing regenerate_key_fn and call the rotation hook """ - # Create regenerate request + # Create regenerate request regenerate_request = RegenerateKeyRequest( - key=key.token or "" + key=key.token or "", + key_alias=key.key_alias, # Pass key alias to ensure correct secret is updated in AWS Secrets Manager ) - + # Create a system user for key rotation from litellm.proxy._types import UserAPIKeyAuth + system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth() - + # Use existing regenerate key function response = await regenerate_key_fn( data=regenerate_request, user_api_key_dict=system_user, - litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, ) - + # Update the NEW key with rotation info (regenerate_key_fn creates a new token) - if isinstance(response, GenerateKeyResponse) and response.token_id and key.rotation_interval: + if ( + isinstance(response, GenerateKeyResponse) + and response.token_id + and key.rotation_interval + ): # Calculate next rotation time using helper function now = datetime.now(timezone.utc) next_rotation_time = _calculate_key_rotation_time(key.rotation_interval) @@ -125,10 +147,10 @@ class KeyRotationManager: data={ "rotation_count": (key.rotation_count or 0) + 1, "last_rotation_at": now, - "key_rotation_at": next_rotation_time - } + "key_rotation_at": next_rotation_time, + }, ) - + # Call the existing rotation hook for notifications, audit logs, etc. if isinstance(response, GenerateKeyResponse): await KeyManagementEventHooks.async_key_rotated_hook( @@ -136,6 +158,5 @@ class KeyRotationManager: existing_key_row=key, response=response, user_api_key_dict=system_user, - litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, ) - \ No newline at end of file diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 9263bca100c..50f8b2a3ded 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -152,7 +152,8 @@ class KeyManagementEventHooks: ) await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager( current_secret_name=initial_secret_name, - new_secret_name=data.key_alias + new_secret_name=response.key_alias + or data.key_alias or f"virtual-key-{response.token_id}", new_secret_value=response.key, ) diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py new file mode 100644 index 00000000000..308c8cdbce1 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py @@ -0,0 +1,144 @@ +""" +Regression test for AWS Secrets Manager Auto-Rotation Bug Fix + +This test verifies that KeyRotationManager correctly passes key_alias +when calling regenerate_key_fn, ensuring the secret is rotated at the +correct location in AWS Secrets Manager. + +Bug Fixed: Key alias was not passed during auto-rotation, causing +secrets to be created at a new location instead of updating in-place. +""" +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + GenerateKeyResponse, + LiteLLM_VerificationToken, + RegenerateKeyRequest, +) +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestKeyRotationManagerPassesKeyAlias: + """ + Regression tests to ensure KeyRotationManager passes key_alias + to regenerate_key_fn during auto-rotation. + """ + + @pytest.mark.asyncio + async def test_rotate_key_passes_key_alias_to_regenerate_request(self): + """ + Verify that _rotate_key includes key_alias in the RegenerateKeyRequest. + + This is the core fix: previously, key_alias was NOT passed, causing + the secret manager hook to use a generated name instead of the alias. + """ + # Create a mock key with an alias + test_alias = "tenant1/my-important-key" + test_token = "sk-test-token-hash-12345" + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_token + mock_key.key_alias = test_alias + mock_key.key_name = "sk-...1234" + mock_key.rotation_interval = "30d" + mock_key.rotation_count = 0 + + # Create mock prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.update = AsyncMock( + return_value=mock_key + ) + + # Create mock response + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + key_alias=test_alias, + ) + + # Capture the RegenerateKeyRequest passed to regenerate_key_fn + captured_request = None + + async def capture_regenerate_key_fn( + data, user_api_key_dict, litellm_changed_by + ): + nonlocal captured_request + captured_request = data + return mock_response + + # Patch regenerate_key_fn to capture the request + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + side_effect=capture_regenerate_key_fn, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + rotation_manager = KeyRotationManager(mock_prisma) + await rotation_manager._rotate_key(mock_key) + + # CRITICAL ASSERTION: key_alias must be passed + assert captured_request is not None, "regenerate_key_fn should have been called" + assert isinstance(captured_request, RegenerateKeyRequest) + assert captured_request.key == test_token, "Token should be passed correctly" + assert captured_request.key_alias == test_alias, ( + f"key_alias should be '{test_alias}' but was '{captured_request.key_alias}'. " + "This is the bug we fixed - key_alias was not being passed!" + ) + + @pytest.mark.asyncio + async def test_rotate_key_passes_none_alias_when_key_has_no_alias(self): + """ + Verify that _rotate_key handles keys without an alias gracefully. + """ + test_token = "sk-test-token-hash-67890" + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_token + mock_key.key_alias = None # No alias set + mock_key.key_name = "sk-...5678" + mock_key.rotation_interval = "30d" + mock_key.rotation_count = 0 + + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.update = AsyncMock( + return_value=mock_key + ) + + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + ) + + captured_request = None + + async def capture_regenerate_key_fn( + data, user_api_key_dict, litellm_changed_by + ): + nonlocal captured_request + captured_request = data + return mock_response + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + side_effect=capture_regenerate_key_fn, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + rotation_manager = KeyRotationManager(mock_prisma) + await rotation_manager._rotate_key(mock_key) + + assert captured_request is not None + assert captured_request.key == test_token + assert ( + captured_request.key_alias is None + ), "key_alias should be None for keys without alias"