phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) Fix - don't allow viewer roles to create virtual keys (#6764)

Browse source 
* fix ui route permissions

* fix test_is_ui_route_allowed

* fix test_is_ui_route_allowed

* test_user_role_permissions
This commit is contained in:
Ishaan Jaff 2024-11-15 18:02:13 -08:00 • committed by GitHub
parent 9ba8f40bd1
commit 6d2ee70a9a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 148 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

91
tests/proxy_admin_ui_tests/test_role_based_access.py
View file

@ -437,3 +437,94 @@ async def test_org_admin_create_user_team_wrong_org_permissions(prisma_client):
"You do not have the required role to call" in e.message
and org2_id in e.message
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"route, user_role, expected_result",
[
# Proxy Admin checks
("/global/spend/logs", LitellmUserRoles.PROXY_ADMIN, True),
("/key/delete", LitellmUserRoles.PROXY_ADMIN, True),
("/key/generate", LitellmUserRoles.PROXY_ADMIN, True),
("/key/regenerate", LitellmUserRoles.PROXY_ADMIN, True),
# # Internal User checks - allowed routes
("/global/spend/logs", LitellmUserRoles.INTERNAL_USER, True),
("/key/delete", LitellmUserRoles.INTERNAL_USER, True),
("/key/generate", LitellmUserRoles.INTERNAL_USER, True),
("/key/82akk800000000jjsk/regenerate", LitellmUserRoles.INTERNAL_USER, True),
# Internal User Viewer
("/key/generate", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
(
"/key/82akk800000000jjsk/regenerate",
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
False,
),
("/key/delete", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
("/team/new", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
("/team/delete", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
("/team/update", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
# Proxy Admin Viewer
("/global/spend/logs", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, True),
("/key/delete", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
("/key/generate", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
(
"/key/82akk800000000jjsk/regenerate",
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
False,
),
("/team/new", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
("/team/delete", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
("/team/update", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
# Internal User checks - disallowed routes
("/organization/member_add", LitellmUserRoles.INTERNAL_USER, False),
],
)
async def test_user_role_permissions(prisma_client, route, user_role, expected_result):
"""Test user role based permissions for different routes"""
try:
# Setup
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
# Admin - admin creates a new user
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="1234",
)
request = NewUserRequest(user_role=user_role)
new_user_response = await new_user(request, user_api_key_dict=user_api_key_dict)
user_id = new_user_response.user_id
# Generate key for new user with team_id="litellm-dashboard"
key_response = await generate_key_fn(
data=GenerateKeyRequest(user_id=user_id, team_id="litellm-dashboard"),
user_api_key_dict=user_api_key_dict,
)
generated_key = key_response.key
bearer_token = "Bearer " + generated_key
# Create request with route
request = Request(scope={"type": "http"})
request._url = URL(url=route)
# Test authorization
if expected_result is True:
# Should pass without error
result = await user_api_key_auth(request=request, api_key=bearer_token)
print(f"Auth passed as expected for {route} with role {user_role}")
else:
# Should raise an error
with pytest.raises(Exception) as exc_info:
await user_api_key_auth(request=request, api_key=bearer_token)
print(f"Auth failed as expected for {route} with role {user_role}")
print(f"Error message: {str(exc_info.value)}")
except Exception as e:
if expected_result:
pytest.fail(f"Expected success but got exception: {str(e)}")
else:
print(f"Got expected exception: {str(e)}")