forked from phoenix/litellm-mirror
(fix) Fix - don't allow viewer roles to create virtual keys (#6764)
* 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
GitHub
parent
9ba8f40bd1
commit
6d2ee70a9a
6 changed files with 148 additions and 42 deletions
|
|
@ -305,14 +305,14 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
|
|||
[
|
||||
# Proxy Admin checks
|
||||
("/global/spend/logs", "proxy_admin", True),
|
||||
("/key/delete", "proxy_admin", True),
|
||||
("/key/generate", "proxy_admin", True),
|
||||
("/key/regenerate", "proxy_admin", True),
|
||||
("/key/delete", "proxy_admin", False),
|
||||
("/key/generate", "proxy_admin", False),
|
||||
("/key/regenerate", "proxy_admin", False),
|
||||
# Internal User checks - allowed routes
|
||||
("/global/spend/logs", "internal_user", True),
|
||||
("/key/delete", "internal_user", True),
|
||||
("/key/generate", "internal_user", True),
|
||||
("/key/82akk800000000jjsk/regenerate", "internal_user", True),
|
||||
("/key/delete", "internal_user", False),
|
||||
("/key/generate", "internal_user", False),
|
||||
("/key/82akk800000000jjsk/regenerate", "internal_user", False),
|
||||
# Internal User Viewer
|
||||
("/key/generate", "internal_user_viewer", False),
|
||||
# Internal User checks - disallowed routes
|
||||
|
|
@ -320,7 +320,7 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
|
|||
],
|
||||
)
|
||||
def test_is_ui_route_allowed(route, user_role, expected_result):
|
||||
from litellm.proxy.auth.user_api_key_auth import _is_ui_route_allowed
|
||||
from litellm.proxy.auth.user_api_key_auth import _is_ui_route
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
|
||||
user_obj = LiteLLM_UserTable(
|
||||
|
|
@ -342,7 +342,7 @@ def test_is_ui_route_allowed(route, user_role, expected_result):
|
|||
"user_obj": user_obj,
|
||||
}
|
||||
try:
|
||||
assert _is_ui_route_allowed(**received_args) == expected_result
|
||||
assert _is_ui_route(**received_args) == expected_result
|
||||
except Exception as e:
|
||||
# If expected result is False, we expect an error
|
||||
if expected_result is False:
|
||||
|
|
|
|||
|
|
@ -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)}")
|
||||
|
|
|
|||
|
|
@ -305,14 +305,14 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
|
|||
[
|
||||
# Proxy Admin checks
|
||||
("/global/spend/logs", "proxy_admin", True),
|
||||
("/key/delete", "proxy_admin", True),
|
||||
("/key/generate", "proxy_admin", True),
|
||||
("/key/regenerate", "proxy_admin", True),
|
||||
("/key/delete", "proxy_admin", False),
|
||||
("/key/generate", "proxy_admin", False),
|
||||
("/key/regenerate", "proxy_admin", False),
|
||||
# Internal User checks - allowed routes
|
||||
("/global/spend/logs", "internal_user", True),
|
||||
("/key/delete", "internal_user", True),
|
||||
("/key/generate", "internal_user", True),
|
||||
("/key/82akk800000000jjsk/regenerate", "internal_user", True),
|
||||
("/key/delete", "internal_user", False),
|
||||
("/key/generate", "internal_user", False),
|
||||
("/key/82akk800000000jjsk/regenerate", "internal_user", False),
|
||||
# Internal User Viewer
|
||||
("/key/generate", "internal_user_viewer", False),
|
||||
# Internal User checks - disallowed routes
|
||||
|
|
@ -320,7 +320,7 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
|
|||
],
|
||||
)
|
||||
def test_is_ui_route_allowed(route, user_role, expected_result):
|
||||
from litellm.proxy.auth.user_api_key_auth import _is_ui_route_allowed
|
||||
from litellm.proxy.auth.user_api_key_auth import _is_ui_route
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
|
||||
user_obj = LiteLLM_UserTable(
|
||||
|
|
@ -342,7 +342,7 @@ def test_is_ui_route_allowed(route, user_role, expected_result):
|
|||
"user_obj": user_obj,
|
||||
}
|
||||
try:
|
||||
assert _is_ui_route_allowed(**received_args) == expected_result
|
||||
assert _is_ui_route(**received_args) == expected_result
|
||||
except Exception as e:
|
||||
# If expected result is False, we expect an error
|
||||
if expected_result is False:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue