mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
(UI) - Security Improvement, move to JWT Auth for Admin UI Sessions (#8995)
* (UI) - Improvements to session handling logic (#8970)
* add cookieUtils
* use utils for clearing cookies
* on logout use clearTokenCookies
* ui use correct clearTokenCookies
* navbar show userEmail on UserID page
* add timestamp on token cookie
* update generate_authenticated_redirect_response
* use common getAuthToken
* fix clearTokenCookies
* fixes for get auth token
* fix invitation link sign in logic
* Revert "fix invitation link sign in logic"
This reverts commit 30e5308cb3.
* fix getAuthToken
* update setAuthToken
* fix ui session handling
* fix ui session handler
* bug fix stop generating LiteLLM Virtual keys for access
* working JWT insert into cookies
* use central place to build UI JWT token
* add _validate_ui_token
* fix ui session handler
* fix fetchWithCredentials
* check allowed routes for ui session tokens
* expose validate_session endpoint
* validate session endpoint
* call sso/session/validate
* getUISessionDetails
* ui move to getUISessionDetails
* /sso/session/validate
* fix cookie utils
* use getUISessionDetails
* use ui_session_id
* "/spend/logs/ui" in spend_tracking_routes
* working sign in JWT flow for proxy admin
* allow proxy admin to access ui routes
* use check_route_access
* update types
* update login method
* fixes to ui session handler
* working flow for admin and internal users
* fixes for invite links
* use JWTs for SSO sign in
* fix /invitation/new flow
* fix code quality checks
* fix _get_ui_session_token_from_cookies
* /organization/list
* ui sso sign in
* TestUISessionHandler
* TestUISessionHandler
This commit is contained in:
Ishaan Jaff
GitHub
parent
42931638df
commit
01a44a4e47
17 changed files with 1104 additions and 538 deletions
|
|
@ -8,7 +8,7 @@ Has all /sso/* routes
|
|||
import asyncio
|
||||
import os
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
|
||||
from typing import TYPE_CHECKING, List, Optional, Union, cast
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
|
||||
|
|
@ -43,6 +43,7 @@ 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:
|
||||
|
|
@ -408,11 +409,7 @@ 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,
|
||||
|
|
@ -426,6 +423,7 @@ 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(
|
||||
|
|
@ -531,16 +529,7 @@ 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:
|
||||
|
|
@ -559,7 +548,6 @@ 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:
|
||||
|
|
@ -632,24 +620,14 @@ 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.value
|
||||
user_role = user_role or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
|
||||
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.value
|
||||
user_role = LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"user_role: {user_role}; ui_access_mode: {ui_access_mode}"
|
||||
|
|
@ -669,30 +647,20 @@ async def auth_callback(request: Request): # noqa: PLR0915
|
|||
disabled_non_admin_personal_key_creation = (
|
||||
get_disabled_non_admin_personal_key_creation()
|
||||
)
|
||||
|
||||
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",
|
||||
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",
|
||||
)
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
litellm_dashboard_ui += "?userID=" + user_id
|
||||
redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
|
||||
redirect_response.set_cookie(key="token", value=jwt_token, secure=True)
|
||||
return redirect_response
|
||||
|
||||
return UISessionHandler.generate_authenticated_redirect_response(
|
||||
redirect_url=litellm_dashboard_ui, jwt_token=jwt_token
|
||||
)
|
||||
|
||||
|
||||
async def insert_sso_user(
|
||||
|
|
@ -778,3 +746,25 @@ 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}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue