diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 0e811669d1a..ad2ad0a5fe5 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -4,10 +4,48 @@ Types for the management endpoints Might include fastapi/proxy requirements.txt related imports """ -from typing import List +from typing import List, Optional, cast from fastapi_sso.sso.base import OpenID +from litellm.proxy._types import LitellmUserRoles + + +def is_valid_litellm_user_role(role_str: str) -> bool: + """ + Check if a string is a valid LitellmUserRoles enum value (case-insensitive). + + Args: + role_str: String to validate (e.g., "proxy_admin", "PROXY_ADMIN", "internal_user") + + Returns: + True if the string matches a valid LitellmUserRoles value, False otherwise + """ + try: + # Use _value2member_map_ for O(1) lookup, case-insensitive + return role_str.lower() in LitellmUserRoles._value2member_map_ + except Exception: + return False + + +def get_litellm_user_role(role_str: str) -> Optional[LitellmUserRoles]: + """ + Convert a string to a LitellmUserRoles enum if valid (case-insensitive). + + Args: + role_str: String to convert (e.g., "proxy_admin", "PROXY_ADMIN", "internal_user") + + Returns: + LitellmUserRoles enum if valid, None otherwise + """ + try: + # Use _value2member_map_ for O(1) lookup, case-insensitive + result = LitellmUserRoles._value2member_map_.get(role_str.lower()) + return cast(Optional[LitellmUserRoles], result) + except Exception: + return None + class CustomOpenID(OpenID): team_ids: List[str] + user_role: Optional[LitellmUserRoles] = None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 9eec551fb7a..d8f426e1ae6 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -58,7 +58,7 @@ from litellm.proxy.management_endpoints.sso_helper_utils import ( has_admin_ui_access, ) from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add -from litellm.proxy.management_endpoints.types import CustomOpenID +from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role from litellm.proxy.utils import ( PrismaClient, ProxyLogging, @@ -277,6 +277,7 @@ def generic_response_convertor( last_name=response.get(generic_user_last_name_attribute_name), provider=response.get(generic_provider_attribute_name), team_ids=all_teams, + user_role=None, ) @@ -1145,7 +1146,7 @@ class SSOAuthenticationHandler: ) -> str: """ Get the redirect URL for SSO - + Note: existing_key is not added to the URL to avoid changing the callback URL. It should be passed via the state parameter instead. """ @@ -1348,7 +1349,7 @@ class SSOAuthenticationHandler: Checks the request 'source' if a cli state token was passed in This is used to authenticate through the CLI login flow. - + The state parameter format is: {PREFIX}:{key}:{existing_key} - If existing_key is provided, it's included in the state - The state parameter is used to pass data through the OAuth flow without changing the callback URL @@ -1673,22 +1674,49 @@ class MicrosoftSSOHandler: access_token=microsoft_sso.access_token ) + # Extract app roles from the id_token JWT + app_roles = MicrosoftSSOHandler.get_app_roles_from_id_token( + id_token=microsoft_sso.id_token + ) + verbose_proxy_logger.debug(f"Extracted app roles from id_token: {app_roles}") + + # Combine groups and app roles + user_role: Optional[LitellmUserRoles] = None + if app_roles: + # Check if any app role is a valid LitellmUserRoles + for role_str in app_roles: + role = get_litellm_user_role(role_str) + if role is not None: + user_role = role + verbose_proxy_logger.debug( + f"Found valid LitellmUserRoles '{role.value}' in app_roles" + ) + break + + verbose_proxy_logger.debug( + f"Combined team_ids (groups + app roles): {user_team_ids}" + ) + # if user is trying to get the raw sso response for debugging, return the raw sso response if return_raw_sso_response: original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = ( user_team_ids ) + original_msft_result["app_roles"] = app_roles return original_msft_result or {} result = MicrosoftSSOHandler.openid_from_response( response=original_msft_result, team_ids=user_team_ids, + user_role=user_role, ) return result @staticmethod def openid_from_response( - response: Optional[dict], team_ids: List[str] + response: Optional[dict], + team_ids: List[str], + user_role: Optional[LitellmUserRoles], ) -> CustomOpenID: response = response or {} verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}") @@ -1700,10 +1728,54 @@ class MicrosoftSSOHandler: first_name=response.get("givenName"), last_name=response.get("surname"), team_ids=team_ids, + user_role=user_role, ) verbose_proxy_logger.debug(f"Microsoft SSO OpenID Response: {openid_response}") return openid_response + @staticmethod + def get_app_roles_from_id_token(id_token: Optional[str]) -> List[str]: + """ + Extract app roles from the Microsoft Entra ID (Azure AD) id_token JWT. + + App roles are assigned in the Azure AD Enterprise Application and appear + in the 'roles' claim of the id_token. + + Args: + id_token (Optional[str]): The JWT id_token from Microsoft SSO + + Returns: + List[str]: List of app role names assigned to the user + """ + if not id_token: + verbose_proxy_logger.debug("No id_token provided for app role extraction") + return [] + + try: + import jwt + + # Decode the JWT without signature verification + # (signature is already verified by fastapi_sso) + decoded_token = jwt.decode(id_token, options={"verify_signature": False}) + + # Extract roles claim from the token + roles = decoded_token.get("roles", []) + + if roles and isinstance(roles, list): + verbose_proxy_logger.debug( + f"Found {len(roles)} app role(s) in id_token: {roles}" + ) + return roles + else: + verbose_proxy_logger.debug( + "No app roles found in id_token or roles claim is not a list" + ) + return [] + + except Exception as e: + verbose_proxy_logger.error(f"Error extracting app roles from id_token: {e}") + return [] + @staticmethod async def get_user_groups_from_graph_api( access_token: Optional[str] = None,