From 41732696c6595340e4dc366e71196a5da0477cb6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 18 Dec 2025 13:25:11 -0800 Subject: [PATCH] replicate delete checks for regenerate --- .../key_management_endpoints.py | 42 +- .../test_key_management_endpoints.py | 401 +++++++++++++++++- 2 files changed, 424 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2e5a92ca01d..67dce22b535 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2310,24 +2310,36 @@ async def _team_key_deletion_check( return False -async def can_delete_verification_token( +async def can_modify_verification_token( key_info: LiteLLM_VerificationToken, user_api_key_cache: DualCache, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, ) -> bool: """ - - check if user is proxy admin - - check if user is team admin and key is a team key - - check if key is personal key + Check if user has permission to modify (delete/regenerate) a verification token. + + Rules: + - Proxy admin can modify any key + - For team keys: only team admin or key owner can modify + - For personal keys: only key owner can modify + + Args: + key_info: The verification token to check + user_api_key_cache: Cache for user API keys + user_api_key_dict: The user making the request + prisma_client: Prisma client for database access + + Returns: + True if user can modify the key, False otherwise """ is_team_key = _is_team_key(data=key_info) - # 1. Proxy admin can delete any key + # 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 delete + # 2. 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( @@ -2354,7 +2366,7 @@ async def can_delete_verification_token( # Not team admin and doesn't own the key return False - # 3. For personal keys: only key owner can delete + # 3. 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 @@ -2362,6 +2374,8 @@ async def can_delete_verification_token( return False + + async def delete_verification_tokens( tokens: List, user_api_key_cache: DualCache, @@ -2413,7 +2427,7 @@ async def delete_verification_tokens( for key in _keys_being_deleted: async def _delete_key(key: LiteLLM_VerificationToken): - if await can_delete_verification_token( + if await can_modify_verification_token( key_info=key, user_api_key_cache=user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2764,6 +2778,18 @@ async def regenerate_key_fn( user_api_key_cache=user_api_key_cache, ) + # check if user has ownership permission to regenerate key + if not await can_modify_verification_token( + key_info=_key_in_db, + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "You are not authorized to regenerate this key"}, + ) + verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) new_token = get_new_token(data=data) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index b84eaf45922..ff85e6d9e73 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -30,7 +30,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_team_key_limits, _common_key_generation_helper, _list_key_helper, - can_delete_verification_token, + can_modify_verification_token, check_org_key_model_specific_limits, check_team_key_model_specific_limits, generate_key_helper_fn, @@ -2635,7 +2635,7 @@ async def test_can_delete_verification_token_proxy_admin_team_key(monkeypatch): mock_prisma_client = AsyncMock() mock_user_api_key_cache = MagicMock() - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2663,7 +2663,7 @@ async def test_can_delete_verification_token_proxy_admin_personal_key(monkeypatc mock_prisma_client = AsyncMock() mock_user_api_key_cache = MagicMock() - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2714,7 +2714,7 @@ async def test_can_delete_verification_token_team_admin_own_team(monkeypatch): mock_get_team_object, ) - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2765,7 +2765,7 @@ async def test_can_delete_verification_token_team_admin_different_team(monkeypat mock_get_team_object, ) - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2815,7 +2815,7 @@ async def test_can_delete_verification_token_key_owner_team_key(monkeypatch): mock_get_team_object, ) - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2843,7 +2843,7 @@ async def test_can_delete_verification_token_key_owner_personal_key(monkeypatch) mock_prisma_client = AsyncMock() mock_user_api_key_cache = MagicMock() - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2895,7 +2895,7 @@ async def test_can_delete_verification_token_other_user_team_key(monkeypatch): mock_get_team_object, ) - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2923,7 +2923,7 @@ async def test_can_delete_verification_token_other_user_personal_key(monkeypatch mock_prisma_client = AsyncMock() mock_user_api_key_cache = MagicMock() - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2959,7 +2959,7 @@ async def test_can_delete_verification_token_team_key_no_team_found(monkeypatch) mock_get_team_object, ) - result = await can_delete_verification_token( + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict, @@ -2987,7 +2987,386 @@ async def test_can_delete_verification_token_personal_key_no_user_id(monkeypatch mock_prisma_client = AsyncMock() mock_user_api_key_cache = MagicMock() - result = await can_delete_verification_token( + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + +@pytest.mark.asyncio +async def test_can_modify_verification_token_proxy_admin_team_key(monkeypatch): + """Test that proxy admin can modify any team key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + api_key="sk-admin", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_proxy_admin_personal_key(monkeypatch): + """Test that proxy admin can modify any personal key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + api_key="sk-admin", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_team_admin_own_team(monkeypatch): + """Test that team admin can modify team keys from their own team.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="team-admin-user", + api_key="sk-user", + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=None, + rpm_limit=None, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[ + Member(user_id="team-admin-user", role="admin"), + Member(user_id="other-user", role="user"), + ], + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return team_table + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_team_admin_different_team(monkeypatch): + """Test that team admin cannot modify team keys from a different team.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id="test-team-456", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="team-admin-user", + api_key="sk-user", + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-456", + team_alias="test-team", + tpm_limit=None, + rpm_limit=None, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[ + Member(user_id="different-admin", role="admin"), + Member(user_id="other-user", role="user"), + ], + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return team_table + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_key_owner_team_key(monkeypatch): + """Test that key owner can modify their own team key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="key-owner-user", + api_key="sk-user", + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=None, + rpm_limit=None, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[ + Member(user_id="key-owner-user", role="user"), + ], + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return team_table + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_key_owner_personal_key(monkeypatch): + """Test that key owner can modify their own personal key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="key-owner-user", + api_key="sk-user", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_other_user_team_key(monkeypatch): + """Test that other user cannot modify team keys they don't own and aren't admin for.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="other-user", + api_key="sk-user", + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=None, + rpm_limit=None, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[ + Member(user_id="key-owner-user", role="user"), + Member(user_id="other-user", role="user"), + Member(user_id="team-admin-user", role="admin"), + ], + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return team_table + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_other_user_personal_key(monkeypatch): + """Test that other user cannot modify personal keys they don't own.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="other-user", + api_key="sk-user", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_team_key_no_team_found(monkeypatch): + """Test that modification fails when team is not found in database.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id="non-existent-team", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="key-owner-user", + api_key="sk-user", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + + +@pytest.mark.asyncio +async def test_can_modify_verification_token_personal_key_no_user_id(monkeypatch): + """Test that modification fails for personal key when key has no user_id.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id=None, + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="some-user", + api_key="sk-user", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_modify_verification_token( key_info=key_info, user_api_key_cache=mock_user_api_key_cache, user_api_key_dict=user_api_key_dict,