diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index ee38eb6515..451f110a89 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -51,6 +51,7 @@ from litellm.proxy.common_utils.admin_ui_utils import ( from litellm.proxy.common_utils.html_forms.jwt_display_template import ( jwt_display_template, ) +from litellm.proxy.common_utils.html_forms.ui_login import html_form from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso_helper_utils import ( check_is_admin_only_access, @@ -76,20 +77,16 @@ router = APIRouter() @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) -async def serve_login_page( - request: Request, - source: Optional[str] = None, - key: Optional[str] = None, - error: Optional[str] = None, -): +async def google_login(request: Request, source: Optional[str] = None, key: Optional[str] = None): # noqa: PLR0915 """ Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" Example: - Serves a unified login page with options for both normal - username/password login and SSO. """ - from litellm.proxy.proxy_server import premium_user + from litellm.proxy.proxy_server import ( + premium_user, + user_custom_ui_sso_sign_in_handler, + ) microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None) google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) @@ -102,334 +99,6 @@ async def serve_login_page( if is_disabled: return admin_ui_disabled() - ####### Check if user is a Enterprise / Premium User for SSO ####### - sso_available = False - if ( - microsoft_client_id is not None - or google_client_id is not None - or generic_client_id is not None - ): - if premium_user is True: - sso_available = True - - ####### Detect DB + MASTER KEY in .env ####### - missing_env_vars = show_missing_vars_in_env() - if missing_env_vars is not None: - return missing_env_vars - ######################################################### - # Construct Redirect URL - base_url_to_redirect_to: Optional[str] = None - base_url_to_redirect_to = os.getenv("PROXY_BASE_URL", "") - server_root_path = os.getenv("SERVER_ROOT_PATH", "") - if server_root_path != "": - base_url_to_redirect_to += server_root_path - ######################################################### - - # Build the unified login page HTML - error_message = "" - if error == "1": - error_message = """ -
- ⚠️ Invalid username or password. Please try again. -
- """ - - sso_button = "" - if sso_available: - sso_login_url = base_url_to_redirect_to - if sso_login_url.endswith("/"): - sso_login_url += "sso/login" - else: - sso_login_url += "/sso/login" - - sso_button = f""" -
-

or

- - 🔐 Login with SSO - -
- """ - - if base_url_to_redirect_to.endswith("/"): - url_to_redirect_to = base_url_to_redirect_to + "login" - else: - url_to_redirect_to = base_url_to_redirect_to + "/login" - - unified_login_html = f""" - - - - - LiteLLM Login - - - - -
-
- -
-

Login

-

Access your LiteLLM Admin UI.

- - {error_message} - -
-
- - - - - - Default Credentials -
-

By default, Username is admin and Password is your set LiteLLM Proxy MASTER_KEY.

-

Need to set UI credentials or SSO? Check the documentation.

-
- - - - - - -
- - -
- - - {sso_button} -
- - - - """ - - from fastapi.responses import HTMLResponse - - return HTMLResponse(content=unified_login_html, status_code=200) - - -@router.get("/sso/login", tags=["experimental"], include_in_schema=False) -async def sso_login_redirect( - request: Request, source: Optional[str] = None, key: Optional[str] = None -): - """ - Handles SSO login redirect - this is what the "Login with SSO" button points to - """ - from litellm.proxy.proxy_server import ( - premium_user, - user_custom_ui_sso_sign_in_handler, - ) - - microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None) - google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) - generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) - ####### Check if user is a Enterprise / Premium User ####### if ( microsoft_client_id is not None @@ -444,12 +113,18 @@ async def sso_login_redirect( code=status.HTTP_403_FORBIDDEN, ) + ####### Detect DB + MASTER KEY in .env ####### + missing_env_vars = show_missing_vars_in_env() + if missing_env_vars is not None: + return missing_env_vars + ui_username = os.getenv("UI_USERNAME") + # get url from request - always use regular callback, but set state for CLI redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( request=request, sso_callback_route="sso/callback", ) - + # Store CLI key in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( source=source, @@ -462,14 +137,11 @@ async def sso_login_redirect( from litellm_enterprise.proxy.auth.custom_sso_handler import ( EnterpriseCustomSSOHandler, ) - return await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( request=request, ) except ImportError: - raise ValueError( - "Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise." - ) + raise ValueError("Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise.") # Check if we should use SSO handler if ( @@ -488,9 +160,16 @@ async def sso_login_redirect( generic_client_id=generic_client_id, state=cli_state, ) + elif ui_username is not None: + # No Google, Microsoft SSO + # Use UI Credentials set in .env + from fastapi.responses import HTMLResponse + + return HTMLResponse(content=html_form, status_code=200) else: - # No SSO configured, redirect back to login page - return RedirectResponse(url="/sso/key/generate", status_code=303) + from fastapi.responses import HTMLResponse + + return HTMLResponse(content=html_form, status_code=200) def generic_response_convertor( @@ -846,16 +525,15 @@ async def check_and_update_if_proxy_admin_id( async def auth_callback(request: Request, state: Optional[str] = None): # noqa: PLR0915 """Verify login""" verbose_proxy_logger.info(f"Starting SSO callback with state: {state}") - + # Check if this is a CLI login (state starts with our CLI prefix) from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX - if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): # Extract the key ID from the state key_id = state.split(":", 1)[1] verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}") return await cli_sso_callback(request, key=key_id) - + from litellm.proxy._types import LiteLLM_JWTAuth from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.proxy_server import ( @@ -930,7 +608,7 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: status_code=401, detail="Result not returned by SSO provider.", ) - + return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=result, request=request, @@ -940,26 +618,28 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: ) + + async def cli_sso_callback(request: Request, key: Optional[str] = None): """CLI SSO callback - generates the key with pre-specified ID""" verbose_proxy_logger.info(f"CLI SSO callback for key: {key}") - + from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, ) from litellm.proxy.proxy_server import prisma_client - - if not key or not key.startswith("sk-"): + + if not key or not key.startswith('sk-'): raise HTTPException( status_code=400, - detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'", + detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'" ) - + if prisma_client is None: raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) - + # Generate a simple key for CLI usage with the pre-specified key ID try: await generate_key_helper_fn( @@ -973,57 +653,63 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None): table_name="key", token=key, # Use the pre-specified key ID ) - + verbose_proxy_logger.info(f"Generated CLI key: {key}") - + # Return success page from fastapi.responses import HTMLResponse from litellm.proxy.common_utils.html_forms.cli_sso_success import ( render_cli_sso_success_page, ) - + html_content = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) - + except Exception as e: verbose_proxy_logger.error(f"Error generating CLI key: {e}") - raise HTTPException(status_code=500, detail=f"Failed to generate key: {str(e)}") + raise HTTPException( + status_code=500, + detail=f"Failed to generate key: {str(e)}" + ) @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) async def cli_poll_key(key_id: str): """CLI polling endpoint - checks if key exists in DB""" from litellm.proxy.proxy_server import prisma_client - - if not key_id.startswith("sk-"): - raise HTTPException(status_code=400, detail="Invalid key ID format") - + + if not key_id.startswith('sk-'): + raise HTTPException( + status_code=400, + detail="Invalid key ID format" + ) + if prisma_client is None: raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) - + try: # Check if key exists in database from litellm.proxy.utils import hash_token - hashed_token = hash_token(key_id) - + key_obj = await prisma_client.db.litellm_verificationtoken.find_unique( where={"token": hashed_token} ) - + if key_obj: verbose_proxy_logger.info(f"CLI key found: {key_id}") return {"status": "ready", "key": key_id} else: return {"status": "pending"} - + except Exception as e: verbose_proxy_logger.error(f"Error polling for CLI key: {e}") raise HTTPException( - status_code=500, detail=f"Error checking key status: {str(e)}" + status_code=500, + detail=f"Error checking key status: {str(e)}" ) @@ -1125,7 +811,6 @@ class SSOAuthenticationHandler: """ Handler for SSO Authentication across all SSO providers """ - @staticmethod async def get_sso_login_redirect( redirect_url: str, @@ -1478,6 +1163,7 @@ class SSOAuthenticationHandler: _new_team_request.update(_default_team_params) team_request = NewTeamRequest(**_new_team_request) return team_request + @staticmethod def _get_cli_state(source: Optional[str], key: Optional[str]) -> Optional[str]: @@ -1490,15 +1176,13 @@ class SSOAuthenticationHandler: LITELLM_CLI_SESSION_TOKEN_PREFIX, LITELLM_CLI_SOURCE_IDENTIFIER, ) + return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" if source == LITELLM_CLI_SOURCE_IDENTIFIER and key else None + + - return ( - f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" - if source == LITELLM_CLI_SOURCE_IDENTIFIER and key - else None - ) @staticmethod - async def get_redirect_response_from_openid( # noqa: PLR0915 + async def get_redirect_response_from_openid( # noqa: PLR0915 result: Union[OpenID, dict, CustomOpenID], request: Request, received_response: Optional[dict] = None, @@ -1518,18 +1202,14 @@ class SSOAuthenticationHandler: ) from litellm.proxy.utils import get_prisma_client_or_throw from litellm.types.proxy.ui_sso import ReturnedUITokenObject + prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - prisma_client = get_prisma_client_or_throw( - "Prisma client is None, connect a database to your proxy" - ) # User is Authe'd in - generate key for the UI to access Proxy verbose_proxy_logger.info(f"SSO callback result: {result}") user_email: Optional[str] = getattr(result, "email", None) - user_id: Optional[str] = ( - getattr(result, "id", None) if result is not None else None - ) + user_id: Optional[str] = getattr(result, "id", None) if result is not None else None if user_email is not None and os.getenv("ALLOWED_EMAIL_DOMAINS") is not None: email_domain = user_email.split("@")[1] @@ -1714,8 +1394,7 @@ class SSOAuthenticationHandler: redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) redirect_response.set_cookie(key="token", value=jwt_token) return redirect_response - - + class MicrosoftSSOHandler: """ Handles Microsoft SSO callback response and returns a CustomOpenID object @@ -2215,28 +1894,3 @@ async def debug_sso_callback(request: Request): ) return HTMLResponse(content=html_content) - - -@router.post("/sso/key/generate", tags=["experimental"], include_in_schema=False) -async def process_login(request: Request): - """ - Process username/password login from the unified login page - """ - try: - # Get form data - form_data = await request.form() - username = form_data.get("username") - password = form_data.get("password") - - if not username or not password: - return RedirectResponse(url="/sso/key/generate?error=1", status_code=303) - - # Import the actual login function from proxy_server - from litellm.proxy.proxy_server import login - - # Call the real login function that handles all the authentication properly - return await login(request) - - except Exception as e: - verbose_proxy_logger.error(f"Error processing login: {e}") - return RedirectResponse(url="/sso/key/generate?error=1", status_code=303) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 245f350be1..f1565fd55f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -938,10 +938,10 @@ class TestUISSO_FunctionsExistence: from litellm.proxy.management_endpoints.ui_sso import auth_callback assert callable(auth_callback) - def test_sso_login_redirect_exists(self): - """Test that sso_login_redirect function exists""" - from litellm.proxy.management_endpoints.ui_sso import sso_login_redirect - assert callable(sso_login_redirect) + def test_google_login_exists(self): + """Test that google_login function exists""" + from litellm.proxy.management_endpoints.ui_sso import google_login + assert callable(google_login) def test_sso_authentication_handler_exists(self): """Test that SSOAuthenticationHandler class exists with new methods""" @@ -1054,7 +1054,7 @@ class TestCustomUISSO: """Test that proper error is raised when enterprise module is not available""" from unittest.mock import MagicMock, patch - from litellm.proxy.management_endpoints.ui_sso import sso_login_redirect + from litellm.proxy.management_endpoints.ui_sso import google_login # Mock request mock_request = MagicMock()