diff --git a/litellm/proxy/auth_v2/authenticators.py b/litellm/proxy/auth_v2/authenticators.py index dad667a57e3..44cba56332b 100644 --- a/litellm/proxy/auth_v2/authenticators.py +++ b/litellm/proxy/auth_v2/authenticators.py @@ -32,22 +32,15 @@ from .models import ( ) from .network import ip_in_trusted_proxies from .oidc.config import OIDCProviderConfig -from .rbac import Role +from .rbac import filter_claim_roles AT_JWT_TYPES = {"at+jwt", "application/at+jwt"} -_PLATFORM_ROLE_VALUES = {Role.PLATFORM_ADMIN.value, Role.PLATFORM_VIEWER.value} def _apply_role_policy(claims: Dict[str, Any], provider: OIDCProviderConfig) -> None: - raw = claims.get("roles") - if not isinstance(raw, list): - claims["roles"] = [] - return - allowed = set(provider.allowed_roles) - filtered = [role for role in raw if role in allowed] - if not provider.allow_platform_roles: - filtered = [role for role in filtered if role not in _PLATFORM_ROLE_VALUES] - claims["roles"] = filtered + claims["roles"] = filter_claim_roles( + claims.get("roles"), provider.allowed_roles, provider.allow_platform_roles + ) @runtime_checkable diff --git a/litellm/proxy/auth_v2/rbac.py b/litellm/proxy/auth_v2/rbac.py index 31b2cfff907..1924497f9fe 100644 --- a/litellm/proxy/auth_v2/rbac.py +++ b/litellm/proxy/auth_v2/rbac.py @@ -1,7 +1,7 @@ from __future__ import annotations from enum import Enum -from typing import TYPE_CHECKING, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, List, Optional, Tuple import casbin from fastapi.security import SecurityScopes @@ -19,6 +19,21 @@ class Role(str, Enum): TEAM_MEMBER = "team_member" +_PLATFORM_ROLE_VALUES = {Role.PLATFORM_ADMIN.value, Role.PLATFORM_VIEWER.value} + + +def filter_claim_roles( + roles: Any, allowed_roles: List[str], allow_platform_roles: bool +) -> List[str]: + if not isinstance(roles, list): + return [] + allowed = set(allowed_roles) + filtered = [role for role in roles if role in allowed] + if not allow_platform_roles: + filtered = [role for role in filtered if role not in _PLATFORM_ROLE_VALUES] + return filtered + + def has_required_scopes( security_scopes: SecurityScopes, principal: "Principal" ) -> bool: diff --git a/litellm/proxy/auth_v2/saml/config.py b/litellm/proxy/auth_v2/saml/config.py index 4c83fd87225..c2910400245 100644 --- a/litellm/proxy/auth_v2/saml/config.py +++ b/litellm/proxy/auth_v2/saml/config.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional +from typing import Dict, List, Optional from pydantic import BaseModel, Field, model_validator @@ -28,6 +28,8 @@ class SAMLConfig(BaseModel): attribute_map: Dict[str, str] = Field( default_factory=lambda: dict(DEFAULT_SAML_ATTRIBUTE_MAP) ) + allowed_roles: List[str] = Field(default_factory=list) + allow_platform_roles: bool = False @model_validator(mode="after") def _require_idp_metadata(self) -> "SAMLConfig": diff --git a/litellm/proxy/auth_v2/saml/router.py b/litellm/proxy/auth_v2/saml/router.py index 20569d7324b..3a260a98363 100644 --- a/litellm/proxy/auth_v2/saml/router.py +++ b/litellm/proxy/auth_v2/saml/router.py @@ -13,6 +13,7 @@ from scim2_models import Email, Name from scim2_models import User as ScimUser from .config import SAMLConfig +from ..rbac import filter_claim_roles from ..resolver import ProvisioningStore from ..session import safe_relay_state @@ -238,6 +239,9 @@ def build_saml_router(auth: "AuthSecurity") -> APIRouter: name_id = authn_response.get_subject().text ava = authn_response.get_identity() or {} mapped = _map_attributes(ava, config.attribute_map) + mapped["roles"] = filter_claim_roles( + mapped.get("roles"), config.allowed_roles, config.allow_platform_roles + ) user = _user_from_mapped(name_id, mapped) store = cast(ProvisioningStore, auth.resolver)