use static methods for Routechecks

This commit is contained in:
Ishaan Jaff 2024-10-25 10:26:43 +04:00
parent c4cab8812a
commit 2e0f501b56
2 changed files with 163 additions and 160 deletions

View file

@ -17,7 +17,10 @@ from .auth_checks_organization import _user_is_org_admin
from .auth_utils import _has_user_setup_sso from .auth_utils import _has_user_setup_sso
def non_proxy_admin_allowed_routes_check( class RouteChecks:
@staticmethod
def non_proxy_admin_allowed_routes_check(
user_obj: Optional[LiteLLM_UserTable], user_obj: Optional[LiteLLM_UserTable],
_user_role: Optional[LitellmUserRoles], _user_role: Optional[LitellmUserRoles],
route: str, route: str,
@ -25,17 +28,17 @@ def non_proxy_admin_allowed_routes_check(
valid_token: UserAPIKeyAuth, valid_token: UserAPIKeyAuth,
api_key: str, api_key: str,
request_data: dict, request_data: dict,
): ):
""" """
Checks if Non Proxy Admin User is allowed to access the route Checks if Non Proxy Admin User is allowed to access the route
""" """
# Check user has defined custom admin routes # Check user has defined custom admin routes
custom_admin_only_route_check( RouteChecks.custom_admin_only_route_check(
route=route, route=route,
) )
if is_llm_api_route(route=route): if RouteChecks.is_llm_api_route(route=route):
pass pass
elif ( elif (
route in LiteLLMRoutes.info_routes.value route in LiteLLMRoutes.info_routes.value
@ -78,7 +81,7 @@ def non_proxy_admin_allowed_routes_check(
pass pass
elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value:
if is_llm_api_route(route=route): if RouteChecks.is_llm_api_route(route=route):
raise HTTPException( raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, status_code=status.HTTP_403_FORBIDDEN,
detail=f"user not allowed to access this OpenAI routes, role= {_user_role}", detail=f"user not allowed to access this OpenAI routes, role= {_user_role}",
@ -131,8 +134,8 @@ def non_proxy_admin_allowed_routes_check(
f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={user_id}" f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={user_id}"
) )
@staticmethod
def custom_admin_only_route_check(route: str): def custom_admin_only_route_check(route: str):
from litellm.proxy.proxy_server import general_settings, premium_user from litellm.proxy.proxy_server import general_settings, premium_user
if "admin_only_routes" in general_settings: if "admin_only_routes" in general_settings:
@ -148,8 +151,8 @@ def custom_admin_only_route_check(route: str):
) )
pass pass
@staticmethod
def is_llm_api_route(route: str) -> bool: def is_llm_api_route(route: str) -> bool:
""" """
Helper to checks if provided route is an OpenAI route Helper to checks if provided route is an OpenAI route

View file

@ -69,7 +69,7 @@ from litellm.proxy.auth.auth_utils import (
) )
from litellm.proxy.auth.oauth2_check import check_oauth2_token from litellm.proxy.auth.oauth2_check import check_oauth2_token
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
from litellm.proxy.auth.route_checks import non_proxy_admin_allowed_routes_check from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.service_account_checks import service_account_checks from litellm.proxy.auth.service_account_checks import service_account_checks
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.utils import _to_ns from litellm.proxy.utils import _to_ns
@ -150,7 +150,7 @@ def _is_api_route_allowed(
raise Exception("Invalid proxy server token passed") raise Exception("Invalid proxy server token passed")
if not _is_user_proxy_admin(user_obj=user_obj): # if non-admin if not _is_user_proxy_admin(user_obj=user_obj): # if non-admin
non_proxy_admin_allowed_routes_check( RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj, user_obj=user_obj,
_user_role=_user_role, _user_role=_user_role,
route=route, route=route,