diff --git a/litellm/proxy/management_endpoints/sso/__init__.py b/litellm/proxy/management_endpoints/sso/__init__.py new file mode 100644 index 00000000000..8144e7c53ff --- /dev/null +++ b/litellm/proxy/management_endpoints/sso/__init__.py @@ -0,0 +1,12 @@ +""" +SSO (Single Sign-On) related modules for LiteLLM Proxy. + +This package contains custom SSO implementations and utilities. +""" + +from litellm.proxy.management_endpoints.sso.custom_microsoft_sso import ( + CustomMicrosoftSSO, +) + +__all__ = ["CustomMicrosoftSSO"] + diff --git a/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py b/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py new file mode 100644 index 00000000000..61b3a8231ac --- /dev/null +++ b/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py @@ -0,0 +1,91 @@ +""" +Custom Microsoft SSO class that allows overriding default Microsoft endpoints. + +This module provides a subclass of fastapi_sso's MicrosoftSSO that allows +custom authorization, token, and userinfo endpoints to be specified via environment +variables. + +Environment Variables: +- MICROSOFT_AUTHORIZATION_ENDPOINT: Custom authorization endpoint URL +- MICROSOFT_TOKEN_ENDPOINT: Custom token endpoint URL +- MICROSOFT_USERINFO_ENDPOINT: Custom userinfo endpoint URL + +If these are not set, the default Microsoft endpoints are used. +""" + +import os +from typing import List, Optional, Union + +import pydantic +from fastapi_sso.sso.base import DiscoveryDocument +from fastapi_sso.sso.microsoft import MicrosoftSSO + +from litellm._logging import verbose_proxy_logger + + +class CustomMicrosoftSSO(MicrosoftSSO): + """ + Microsoft SSO subclass that allows overriding default endpoints via environment variables. + + Supports: + - MICROSOFT_AUTHORIZATION_ENDPOINT + - MICROSOFT_TOKEN_ENDPOINT + - MICROSOFT_USERINFO_ENDPOINT + """ + + def __init__( + self, + client_id: str, + client_secret: str, + redirect_uri: Optional[Union[pydantic.AnyHttpUrl, str]] = None, + allow_insecure_http: bool = False, + scope: Optional[List[str]] = None, + tenant: Optional[str] = None, + ): + super().__init__( + client_id=client_id, + client_secret=client_secret, + redirect_uri=redirect_uri, + allow_insecure_http=allow_insecure_http, + scope=scope, + tenant=tenant, + ) + + async def get_discovery_document(self) -> DiscoveryDocument: + """ + Override to support custom endpoints via environment variables. + Falls back to default Microsoft endpoints if not set. + """ + custom_authorization_endpoint = os.getenv( + "MICROSOFT_AUTHORIZATION_ENDPOINT", None + ) + custom_token_endpoint = os.getenv("MICROSOFT_TOKEN_ENDPOINT", None) + custom_userinfo_endpoint = os.getenv("MICROSOFT_USERINFO_ENDPOINT", None) + + # Use custom endpoints if set, otherwise use defaults + authorization_endpoint = ( + custom_authorization_endpoint + or f"https://login.microsoftonline.com/{self.tenant}/oauth2/v2.0/authorize" + ) + token_endpoint = ( + custom_token_endpoint + or f"https://login.microsoftonline.com/{self.tenant}/oauth2/v2.0/token" + ) + userinfo_endpoint = ( + custom_userinfo_endpoint or f"https://graph.microsoft.com/{self.version}/me" + ) + + if custom_authorization_endpoint or custom_token_endpoint or custom_userinfo_endpoint: + verbose_proxy_logger.debug( + f"Using custom Microsoft SSO endpoints - " + f"authorization: {authorization_endpoint}, " + f"token: {token_endpoint}, " + f"userinfo: {userinfo_endpoint}" + ) + + return DiscoveryDocument( + authorization_endpoint=authorization_endpoint, + token_endpoint=token_endpoint, + userinfo_endpoint=userinfo_endpoint, + ) +