scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

replicate delete checks for regenerate

Browse source
This commit is contained in:
yuneng-jiang 2025-12-18 13:25:11 -08:00
parent 0a1fb204cd
commit 41732696c6
2 changed files with 424 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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)

401
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -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,