Revert "(UI) - Security Improvement, move to JWT Auth for Admin UI Sessions (#8995)"

This reverts commit 01a44a4e47.

litellm/proxy/management_endpoints/ui_sso.py

@@ -8,7 +8,7 @@ Has all /sso/* routes
 import asyncio
 import os
 import uuid
-from typing import TYPE_CHECKING, List, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 
@@ -43,7 +43,6 @@ from litellm.proxy.management_endpoints.sso_helper_utils import (
 )
 from litellm.proxy.management_endpoints.team_endpoints import team_member_add
 from litellm.proxy.management_endpoints.types import CustomOpenID
-from litellm.proxy.management_helpers.ui_session_handler import UISessionHandler
 from litellm.secret_managers.main import str_to_bool
 
 if TYPE_CHECKING:
@@ -409,7 +408,11 @@ def get_disabled_non_admin_personal_key_creation():
 @router.get("/sso/callback", tags=["experimental"], include_in_schema=False)
 async def auth_callback(request: Request):  # noqa: PLR0915
     """Verify login"""
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        generate_key_helper_fn,
+    )
     from litellm.proxy.proxy_server import (
+        general_settings,
         jwt_handler,
         master_key,
         premium_user,
@@ -423,7 +426,6 @@ async def auth_callback(request: Request):  # noqa: PLR0915
     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)
-    user_role: Optional[LitellmUserRoles] = None
     # get url from request
     if master_key is None:
         raise ProxyException(
@@ -529,7 +531,16 @@ async def auth_callback(request: Request):  # noqa: PLR0915
     max_internal_user_budget = litellm.max_internal_user_budget
     internal_user_budget_duration = litellm.internal_user_budget_duration
 
+    # User might not be already created on first generation of key
     # But if it is, we want their models preferences
+    default_ui_key_values: Dict[str, Any] = {
+        "duration": "24hr",
+        "key_max_budget": litellm.max_ui_session_budget,
+        "aliases": {},
+        "config": {},
+        "spend": 0,
+        "team_id": "litellm-dashboard",
+    }
     user_defined_values: Optional[SSOUserDefinedValues] = None
 
     if user_custom_sso is not None:
@@ -548,6 +559,7 @@ async def auth_callback(request: Request):  # noqa: PLR0915
         )
 
     _user_id_from_sso = user_id
+    user_role = None
     try:
         if prisma_client is not None:
             try:
@@ -620,14 +632,24 @@ async def auth_callback(request: Request):  # noqa: PLR0915
         f"user_defined_values for creating ui key: {user_defined_values}"
     )
 
+    default_ui_key_values.update(user_defined_values)
+    default_ui_key_values["request_type"] = "key"
+    response = await generate_key_helper_fn(
+        **default_ui_key_values,  # type: ignore
+        table_name="key",
+    )
+
+    key = response["token"]  # type: ignore
+    user_id = response["user_id"]  # type: ignore
+
     litellm_dashboard_ui = "/ui/"
-    user_role = user_role or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
+    user_role = user_role or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
     if (
         os.getenv("PROXY_ADMIN_ID", None) is not None
         and os.environ["PROXY_ADMIN_ID"] == user_id
     ):
         # checks if user is admin
-        user_role = LitellmUserRoles.PROXY_ADMIN
+        user_role = LitellmUserRoles.PROXY_ADMIN.value
 
     verbose_proxy_logger.debug(
         f"user_role: {user_role}; ui_access_mode: {ui_access_mode}"
@@ -647,20 +669,30 @@ async def auth_callback(request: Request):  # noqa: PLR0915
     disabled_non_admin_personal_key_creation = (
         get_disabled_non_admin_personal_key_creation()
     )
-    jwt_token = UISessionHandler.build_authenticated_ui_jwt_token(
-        user_id=user_defined_values.get("user_id", ""),
-        user_role=user_role,
-        user_email=user_defined_values.get("user_email", ""),
-        premium_user=premium_user,
-        disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation,
-        login_method="sso",
+
+    import jwt
+
+    jwt_token = jwt.encode(  # type: ignore
+        {
+            "user_id": user_id,
+            "key": key,
+            "user_email": user_email,
+            "user_role": user_role,
+            "login_method": "sso",
+            "premium_user": premium_user,
+            "auth_header_name": general_settings.get(
+                "litellm_key_header_name", "Authorization"
+            ),
+            "disabled_non_admin_personal_key_creation": disabled_non_admin_personal_key_creation,
+        },
+        master_key,
+        algorithm="HS256",
     )
     if user_id is not None and isinstance(user_id, str):
         litellm_dashboard_ui += "?userID=" + user_id
-
-    return UISessionHandler.generate_authenticated_redirect_response(
-        redirect_url=litellm_dashboard_ui, jwt_token=jwt_token
-    )
+    redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
+    redirect_response.set_cookie(key="token", value=jwt_token, secure=True)
+    return redirect_response
 
 
 async def insert_sso_user(
@@ -746,25 +778,3 @@ async def get_ui_settings(request: Request):
         ),
         "DISABLE_EXPENSIVE_DB_QUERIES": disable_expensive_db_queries,
     }
-
-
-@router.get(
-    "/sso/session/validate",
-    include_in_schema=False,
-    tags=["experimental"],
-)
-async def validate_session(
-    request: Request,
-    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
-):
-    from litellm.proxy.auth.handle_jwt import JWTHandler
-    from litellm.proxy.management_helpers.ui_session_handler import UISessionHandler
-
-    ui_session_token = UISessionHandler._get_ui_session_token_from_cookies(request)
-    ui_session_id = UISessionHandler._get_latest_ui_cookie_name(request.cookies)
-    if ui_session_token is None:
-        raise HTTPException(status_code=401, detail="Unauthorized")
-    jwt_handler = JWTHandler()
-    validated_jwt_token = await jwt_handler.auth_jwt(token=ui_session_token)
-    validated_jwt_token["session_id"] = ui_session_id
-    return {"valid": True, "data": validated_jwt_token}