diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 22c4f27d94a..f922cc22f44 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -896,6 +896,29 @@ def determine_role_from_groups( return role_mappings.default_role +def _groups_from_claim(value: object) -> list[str]: + match value: + case None: + return [] + case list(): + groups: Final[Sequence[object]] = value + return [str(group) for group in groups] + case str(): + return [group.strip() for group in value.split(",") if group.strip()] + case _: + return [str(value)] + + +def role_from_role_mappings(claims: Mapping[str, object], role_mappings: RoleMappings) -> LitellmUserRoles | None: + user_groups: Final = _groups_from_claim(get_nested_value(claims, role_mappings.group_claim)) + if not user_groups: + verbose_proxy_logger.debug( + "No groups found in '%s', using default_role: %s", role_mappings.group_claim, role_mappings.default_role + ) + return role_mappings.default_role + return determine_role_from_groups(user_groups, role_mappings) + + def process_sso_jwt_access_token( access_token_str: str | None, sso_jwt_handler: JWTHandler | None, @@ -953,25 +976,7 @@ def process_sso_jwt_access_token( # Try role_mappings first (group-based role determination) if role_mappings is not None and role_mappings.roles: - group_claim: Final = role_mappings.group_claim - user_groups_raw: Final[object] = get_nested_value(access_token_payload, group_claim) - - user_groups: list[str] = [] - if isinstance(user_groups_raw, list): - raw_groups: Final[Sequence[object]] = user_groups_raw - user_groups = [str(g) for g in raw_groups] - elif isinstance(user_groups_raw, str): - user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()] - elif user_groups_raw is not None: - user_groups = [str(user_groups_raw)] - - if user_groups: - user_role = determine_role_from_groups(user_groups, role_mappings) - verbose_proxy_logger.debug( - "Determined role '%s' from access token groups '%s' using role_mappings", user_role, user_groups - ) - elif role_mappings.default_role: - user_role = role_mappings.default_role + user_role = role_from_role_mappings(access_token_payload, role_mappings) # Fallback: try GENERIC_USER_ROLE_ATTRIBUTE on the access token payload if user_role is None: @@ -1268,35 +1273,7 @@ def generic_response_convertor( "generic", "okta", ]: - # Use role_mappings to determine role from groups - group_claim: Final = role_mappings.group_claim - user_groups_raw: Final[object] = get_nested_value(response, group_claim) - - # Handle different formats: could be a list, string (comma-separated), or single value - user_groups: list[str] = [] - if isinstance(user_groups_raw, list): - raw_groups: Final[Sequence[object]] = user_groups_raw - user_groups = [str(g) for g in raw_groups] - elif isinstance(user_groups_raw, str): - # Handle comma-separated string - user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()] - elif user_groups_raw is not None: - # Single value - user_groups = [str(user_groups_raw)] - - if user_groups: - user_role = determine_role_from_groups(user_groups, role_mappings) - verbose_proxy_logger.debug( - "Determined role '%s' from groups '%s' using role_mappings", - user_role.value if user_role else None, - user_groups, - ) - else: - # No groups found, use default_role - user_role = role_mappings.default_role - verbose_proxy_logger.debug( - "No groups found in '%s', using default_role: %s", group_claim, role_mappings.default_role - ) + user_role = role_from_role_mappings(response, role_mappings) # Fallback to existing logic if role_mappings not used if user_role is None: @@ -2165,6 +2142,7 @@ async def _sso_provider_callback( microsoft_client_id=microsoft_client_id, redirect_url=redirect_url, return_raw_sso_response=return_raw_sso_response, + role_mappings=await _setup_role_mappings(), ) return microsoft_result, None, None, microsoft_assertion if generic_client_id is not None: @@ -4575,6 +4553,7 @@ class MicrosoftSSOHandler: microsoft_client_id: str, redirect_url: str, return_raw_sso_response: bool = False, + role_mappings: RoleMappings | None = None, ) -> tuple[CustomOpenID | OpenID | dict[str, object], SSOIdentityAssertion | None]: """ Get the Microsoft SSO callback response, with the IdP identity assertion the token @@ -4630,8 +4609,12 @@ class MicrosoftSSOHandler: app_roles: Final = MicrosoftSSOHandler.get_app_roles_from_id_token(id_token=microsoft_sso.id_token) verbose_proxy_logger.debug("Extracted app roles from id_token: %s", app_roles) - # Combine groups and app roles - user_role: Final = MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles) + mapped_role: Final = ( + role_from_role_mappings(_decode_sso_token_claims(microsoft_tokens.id_token), role_mappings) + if role_mappings is not None and role_mappings.provider.lower() == "microsoft" + else None + ) + user_role: Final = mapped_role or MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles) verbose_proxy_logger.debug("Combined team_ids (groups + app roles): %s", user_team_ids) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 4135e9be6f6..771c8ec1073 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -17,7 +17,7 @@ from pydantic import SecretStr import litellm from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, NewUserResponse, ProxyException from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.types import CustomOpenID @@ -34,6 +34,7 @@ from litellm.proxy.management_endpoints.ui_sso import ( normalize_email, process_sso_jwt_access_token, read_return_to_cookie, + role_from_role_mappings, sso_lookup_user_ids, ) from litellm.types.proxy.management_endpoints.ui_sso import ( @@ -41,6 +42,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( MicrosoftGraphAPIUserGroupDirectoryObject, MicrosoftGraphAPIUserGroupResponse, MicrosoftServicePrincipalTeam, + RoleMappings, TeamMappings, ) from tests._master_key import MASTER_KEY @@ -410,6 +412,99 @@ def test_get_microsoft_callback_response_raw_sso_response(stubbed_graph_api): assert result["surname"] == "User" +@pytest.mark.parametrize( + ("role_mappings", "id_token_claims", "expected_role"), + [ + pytest.param( + RoleMappings(provider="microsoft", group_claim="groups", roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"]}), + {"groups": ["grp-b", "grp-a"]}, + LitellmUserRoles.PROXY_ADMIN, + id="entra-group-mapped", + ), + pytest.param( + RoleMappings( + provider="microsoft", group_claim="roles", roles={LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY: ["Auditor"]} + ), + {"roles": ["Auditor"]}, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + id="entra-app-role-mapped", + ), + pytest.param( + RoleMappings( + provider="microsoft", + group_claim="groups", + default_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"]}, + ), + {}, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + id="no-groups-uses-default-role", + ), + pytest.param( + RoleMappings(provider="microsoft", group_claim="groups", roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"]}), + {"groups": ["grp-b"], "roles": ["internal_user"]}, + LitellmUserRoles.INTERNAL_USER, + id="unmatched-falls-back-to-app-role", + ), + pytest.param( + RoleMappings(provider="generic", group_claim="groups", roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"]}), + {"groups": ["grp-a"]}, + None, + id="other-provider-mappings-ignored", + ), + ], +) +def test_get_microsoft_callback_response_applies_role_mappings( + stubbed_graph_api, role_mappings, id_token_claims, expected_role +): + mock_request = MagicMock(spec=Request) + mock_request.scope = {} + id_token = _ema_id_token(sub="msft123", **id_token_claims) + + async def fake_verify(self, *args, **kwargs): + self._id_token = id_token + return {"mail": "microsoft_user@example.com", "id": "msft123"} + + with patch.dict(os.environ, {"MICROSOFT_CLIENT_SECRET": "mock_secret", "MICROSOFT_TENANT": "mock_tenant"}): + with patch( # test-quality-ok: stubs the upstream Microsoft IdP boundary, where the token response is read + "fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process", + autospec=True, + side_effect=fake_verify, + ): + result, _ = asyncio.run( + MicrosoftSSOHandler.get_microsoft_callback_response( + request=mock_request, + microsoft_client_id="mock_client_id", + redirect_url="http://mock_redirect_url", + role_mappings=role_mappings, + ) + ) + + assert isinstance(result, CustomOpenID) + assert result.user_role == expected_role + + +@pytest.mark.parametrize( + ("claims", "expected_role"), + [ + pytest.param({"groups": ["grp-b", "grp-a"]}, LitellmUserRoles.PROXY_ADMIN, id="list"), + pytest.param({"groups": "grp-b, grp-a"}, LitellmUserRoles.PROXY_ADMIN, id="comma-separated-string"), + pytest.param({"groups": 42}, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, id="single-non-string-value"), + pytest.param({"groups": " , "}, LitellmUserRoles.INTERNAL_USER, id="blank-string-uses-default-role"), + pytest.param({"groups": "grp-z"}, LitellmUserRoles.INTERNAL_USER, id="no-match-uses-default-role"), + ], +) +def test_role_from_role_mappings_reads_every_group_claim_shape(claims, expected_role): + role_mappings = RoleMappings( + provider="microsoft", + group_claim="groups", + default_role=LitellmUserRoles.INTERNAL_USER, + roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"], LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY: ["42"]}, + ) + + assert role_from_role_mappings(claims, role_mappings) == expected_role + + def test_get_google_callback_response(): # Arrange mock_request = MagicMock(spec=Request) @@ -9311,13 +9406,13 @@ async def test_auth_callback_without_oauth_error_proceeds_to_normal_flow(): # ── SSO identity assertion capture + persist wiring (EMA) ───────────────────── -def _ema_id_token(sub: str = "u1") -> str: +def _ema_id_token(sub: str = "u1", **claims: object) -> str: import time as _time import jwt as _pyjwt return _pyjwt.encode( - {"iss": "https://idp.example.com", "sub": sub, "exp": int(_time.time()) + 3600}, + {"iss": "https://idp.example.com", "sub": sub, "exp": int(_time.time()) + 3600, **claims}, "test-idp-signing-key-32-bytes-long-xxxx", algorithm="HS256", )