mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
use static methods for OrganizationRoleBasedAccessChecks
This commit is contained in:
Ishaan Jaff
parent
b3141e1a5f
commit
efbfbbeb30
1 changed files with 167 additions and 129 deletions
|
|
@ -9,153 +9,191 @@ from fastapi import status
|
|||
from litellm.proxy._types import *
|
||||
|
||||
|
||||
def organization_role_based_access_check(
|
||||
request_body: dict,
|
||||
user_object: Optional[LiteLLM_UserTable],
|
||||
route: str,
|
||||
):
|
||||
"""
|
||||
Role based access control checks only run if a user is part of an Organization
|
||||
class OrganizationRoleBasedAccessChecks:
|
||||
|
||||
Organization Checks:
|
||||
ONLY RUN IF user_object.organization_memberships is not None
|
||||
@staticmethod
|
||||
def organization_role_based_access_check(
|
||||
request_body: dict,
|
||||
user_object: Optional[LiteLLM_UserTable],
|
||||
route: str,
|
||||
):
|
||||
"""
|
||||
Role based access control checks only run if a user is part of an Organization
|
||||
|
||||
1. Only Proxy Admins can access /organization/new
|
||||
2. IF route is a LiteLLMRoutes.org_admin_only_routes, then check if user is an Org Admin for that organization
|
||||
Organization Checks:
|
||||
ONLY RUN IF user_object.organization_memberships is not None
|
||||
|
||||
"""
|
||||
1. Only Proxy Admins can access /organization/new
|
||||
2. IF route is a LiteLLMRoutes.org_admin_only_routes, then check if user is an Org Admin for that organization
|
||||
|
||||
if user_object is None:
|
||||
return
|
||||
"""
|
||||
|
||||
passed_organization_id: Optional[str] = request_body.get("organization_id", None)
|
||||
if user_object is None:
|
||||
return
|
||||
|
||||
if route == "/organization/new":
|
||||
if user_object.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||||
raise ProxyException(
|
||||
message=f"Only proxy admins can create new organizations. You are {user_object.user_role}",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="user_role",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
if user_object.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
|
||||
# Checks if route is an Org Admin Only Route
|
||||
if route in LiteLLMRoutes.org_admin_only_routes.value:
|
||||
_user_organizations, _user_organization_role_mapping = (
|
||||
get_user_organization_info(user_object)
|
||||
passed_organization_id: Optional[str] = request_body.get(
|
||||
"organization_id", None
|
||||
)
|
||||
|
||||
if user_object.organization_memberships is None:
|
||||
raise ProxyException(
|
||||
message=f"Tried to access route={route} but you are not a member of any organization. Please contact the proxy admin to request access.",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="organization_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
if passed_organization_id is None:
|
||||
raise ProxyException(
|
||||
message="Passed organization_id is None, please pass an organization_id in your request",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="organization_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
user_role: Optional[LitellmUserRoles] = _user_organization_role_mapping.get(
|
||||
passed_organization_id
|
||||
)
|
||||
if user_role is None:
|
||||
raise ProxyException(
|
||||
message=f"You do not have a role within the selected organization. Passed organization_id: {passed_organization_id}. Please contact the organization admin to request access.",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="organization_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
if user_role != LitellmUserRoles.ORG_ADMIN.value:
|
||||
raise ProxyException(
|
||||
message=f"You do not have the required role to perform {route} in Organization {passed_organization_id}. Your role is {user_role} in Organization {passed_organization_id}",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="user_role",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
elif route == "/team/new":
|
||||
# if user is part of multiple teams, then they need to specify the organization_id
|
||||
_user_organizations, _user_organization_role_mapping = (
|
||||
get_user_organization_info(user_object)
|
||||
)
|
||||
if (
|
||||
user_object.organization_memberships is not None
|
||||
and len(user_object.organization_memberships) > 0
|
||||
):
|
||||
if passed_organization_id is None:
|
||||
if route == "/organization/new":
|
||||
if user_object.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||||
raise ProxyException(
|
||||
message=f"Passed organization_id is None, please specify the organization_id in your request. You are part of multiple organizations: {_user_organizations}",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="organization_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
_user_role_in_passed_org = _user_organization_role_mapping.get(
|
||||
passed_organization_id
|
||||
)
|
||||
if _user_role_in_passed_org != LitellmUserRoles.ORG_ADMIN.value:
|
||||
raise ProxyException(
|
||||
message=f"You do not have the required role to call {route}. Your role is {_user_role_in_passed_org} in Organization {passed_organization_id}",
|
||||
message=f"Only proxy admins can create new organizations. You are {user_object.user_role}",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="user_role",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
if user_object.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
|
||||
def get_user_organization_info(
|
||||
user_object: LiteLLM_UserTable,
|
||||
) -> Tuple[List[str], Dict[str, Optional[LitellmUserRoles]]]:
|
||||
"""
|
||||
Helper function to extract user organization information.
|
||||
# Checks if route is an Org Admin Only Route
|
||||
if route in LiteLLMRoutes.org_admin_only_routes.value:
|
||||
_user_organizations, _user_organization_role_mapping = (
|
||||
OrganizationRoleBasedAccessChecks.get_user_organization_info(
|
||||
user_object
|
||||
)
|
||||
)
|
||||
|
||||
Args:
|
||||
user_object (LiteLLM_UserTable): The user object containing organization memberships.
|
||||
if user_object.organization_memberships is None:
|
||||
raise ProxyException(
|
||||
message=f"Tried to access route={route} but you are not a member of any organization. Please contact the proxy admin to request access.",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="organization_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
Returns:
|
||||
Tuple[List[str], Dict[str, Optional[LitellmUserRoles]]]: A tuple containing:
|
||||
- List of organization IDs the user is a member of
|
||||
- Dictionary mapping organization IDs to user roles
|
||||
"""
|
||||
_user_organizations: List[str] = []
|
||||
_user_organization_role_mapping: Dict[str, Optional[LitellmUserRoles]] = {}
|
||||
if passed_organization_id is None:
|
||||
raise ProxyException(
|
||||
message="Passed organization_id is None, please pass an organization_id in your request",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="organization_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
user_role: Optional[LitellmUserRoles] = _user_organization_role_mapping.get(
|
||||
passed_organization_id
|
||||
)
|
||||
if user_role is None:
|
||||
raise ProxyException(
|
||||
message=f"You do not have a role within the selected organization. Passed organization_id: {passed_organization_id}. Please contact the organization admin to request access.",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="organization_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
if user_role != LitellmUserRoles.ORG_ADMIN.value:
|
||||
raise ProxyException(
|
||||
message=f"You do not have the required role to perform {route} in Organization {passed_organization_id}. Your role is {user_role} in Organization {passed_organization_id}",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="user_role",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
elif route == "/team/new":
|
||||
# if user is part of multiple teams, then they need to specify the organization_id
|
||||
_user_organizations, _user_organization_role_mapping = (
|
||||
OrganizationRoleBasedAccessChecks.get_user_organization_info(
|
||||
user_object
|
||||
)
|
||||
)
|
||||
if (
|
||||
user_object.organization_memberships is not None
|
||||
and len(user_object.organization_memberships) > 0
|
||||
):
|
||||
if passed_organization_id is None:
|
||||
raise ProxyException(
|
||||
message=f"Passed organization_id is None, please specify the organization_id in your request. You are part of multiple organizations: {_user_organizations}",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="organization_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
_user_role_in_passed_org = _user_organization_role_mapping.get(
|
||||
passed_organization_id
|
||||
)
|
||||
if _user_role_in_passed_org != LitellmUserRoles.ORG_ADMIN.value:
|
||||
raise ProxyException(
|
||||
message=f"You do not have the required role to call {route}. Your role is {_user_role_in_passed_org} in Organization {passed_organization_id}",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
param="user_role",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_user_organization_info(
|
||||
user_object: LiteLLM_UserTable,
|
||||
) -> Tuple[List[str], Dict[str, Optional[LitellmUserRoles]]]:
|
||||
"""
|
||||
Helper function to extract user organization information.
|
||||
|
||||
Args:
|
||||
user_object (LiteLLM_UserTable): The user object containing organization memberships.
|
||||
|
||||
Returns:
|
||||
Tuple[List[str], Dict[str, Optional[LitellmUserRoles]]]: A tuple containing:
|
||||
- List of organization IDs the user is a member of
|
||||
- Dictionary mapping organization IDs to user roles
|
||||
"""
|
||||
_user_organizations: List[str] = []
|
||||
_user_organization_role_mapping: Dict[str, Optional[LitellmUserRoles]] = {}
|
||||
|
||||
if user_object.organization_memberships is not None:
|
||||
for _membership in user_object.organization_memberships:
|
||||
if _membership.organization_id is not None:
|
||||
_user_organizations.append(_membership.organization_id)
|
||||
_user_organization_role_mapping[_membership.organization_id] = _membership.user_role # type: ignore
|
||||
|
||||
return _user_organizations, _user_organization_role_mapping
|
||||
|
||||
@staticmethod
|
||||
def _user_is_org_admin_in_requested_org(
|
||||
request_data: dict,
|
||||
user_object: Optional[LiteLLM_UserTable] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Helper function to check if user is an org admin for the passed organization_id in request_data
|
||||
|
||||
Args:
|
||||
request_data (dict): The request data containing the organization_id
|
||||
user_object (LiteLLM_UserTable): The user object containing organization memberships
|
||||
|
||||
Returns:
|
||||
bool: True if user is an org admin for the passed organization_id, False otherwise
|
||||
"""
|
||||
if request_data.get("organization_id", None) is None:
|
||||
return False
|
||||
|
||||
_user_is_org_admin = OrganizationRoleBasedAccessChecks._user_is_admin_in_org(
|
||||
user_object=user_object,
|
||||
organization_id=request_data.get("organization_id", None),
|
||||
)
|
||||
|
||||
return _user_is_org_admin
|
||||
|
||||
@staticmethod
|
||||
def _user_is_admin_in_org(
|
||||
user_object: Optional[LiteLLM_UserTable] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Returns True if user is an admin in the passed organization_id
|
||||
|
||||
Args:
|
||||
user_object (LiteLLM_UserTable): The user object containing organization memberships
|
||||
organization_id (str): The organization_id to check if the user is an admin in
|
||||
|
||||
Returns:
|
||||
bool: True if user is an admin in the passed organization_id, False otherwise
|
||||
"""
|
||||
if user_object is None:
|
||||
return False
|
||||
|
||||
if user_object.organization_memberships is None:
|
||||
return False
|
||||
|
||||
if user_object.organization_memberships is not None:
|
||||
for _membership in user_object.organization_memberships:
|
||||
if _membership.organization_id is not None:
|
||||
_user_organizations.append(_membership.organization_id)
|
||||
_user_organization_role_mapping[_membership.organization_id] = _membership.user_role # type: ignore
|
||||
if _membership.organization_id == organization_id:
|
||||
if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value:
|
||||
return True
|
||||
|
||||
return _user_organizations, _user_organization_role_mapping
|
||||
|
||||
|
||||
def _user_is_org_admin(
|
||||
request_data: dict,
|
||||
user_object: Optional[LiteLLM_UserTable] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Helper function to check if user is an org admin for the passed organization_id
|
||||
"""
|
||||
if request_data.get("organization_id", None) is None:
|
||||
return False
|
||||
|
||||
if user_object is None:
|
||||
return False
|
||||
|
||||
if user_object.organization_memberships is None:
|
||||
return False
|
||||
|
||||
for _membership in user_object.organization_memberships:
|
||||
if _membership.organization_id == request_data.get("organization_id", None):
|
||||
if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue