This commit is contained in:
Timothy J. Baek
2024-10-15 02:48:41 -07:00
parent 2f4c04055c
commit 97bb9d41a6
3 changed files with 62 additions and 15 deletions

View File

@@ -1,10 +1,13 @@
import re
import uuid
import time
import datetime
from open_webui.apps.webui.models.auths import (
AddUserForm,
ApiKey,
Auths,
Token,
SigninForm,
SigninResponse,
SignupForm,
@@ -34,6 +37,7 @@ from open_webui.utils.utils import (
get_password_hash,
)
from open_webui.utils.webhook import post_webhook
from typing import Optional
router = APIRouter()
@@ -42,25 +46,42 @@ router = APIRouter()
############################
@router.get("/", response_model=UserResponse)
class SessionUserResponse(Token, UserResponse):
expires_at: Optional[int] = None
@router.get("/", response_model=SessionUserResponse)
async def get_session_user(
request: Request, response: Response, user=Depends(get_current_user)
):
expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN)
expires_at = None
if expires_delta:
expires_at = int(time.time()) + int(expires_delta.total_seconds())
token = create_token(
data={"id": user.id},
expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
expires_delta=expires_delta,
)
datetime_expires_at = datetime.datetime.fromtimestamp(
expires_at, datetime.timezone.utc
)
# Set the cookie token
response.set_cookie(
key="token",
value=token,
expires=datetime_expires_at,
httponly=True, # Ensures the cookie is not accessible via JavaScript
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
secure=WEBUI_SESSION_COOKIE_SECURE,
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
secure=WEBUI_SESSION_COOKIE_SECURE,
)
return {
"token": token,
"token_type": "Bearer",
"expires_at": expires_at,
"id": user.id,
"email": user.email,
"name": user.name,
@@ -119,7 +140,7 @@ async def update_password(
############################
@router.post("/signin", response_model=SigninResponse)
@router.post("/signin", response_model=SessionUserResponse)
async def signin(request: Request, response: Response, form_data: SigninForm):
if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
if WEBUI_AUTH_TRUSTED_EMAIL_HEADER not in request.headers:
@@ -161,23 +182,35 @@ async def signin(request: Request, response: Response, form_data: SigninForm):
user = Auths.authenticate_user(form_data.email.lower(), form_data.password)
if user:
expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN)
expires_at = None
if expires_delta:
expires_at = int(time.time()) + int(expires_delta.total_seconds())
token = create_token(
data={"id": user.id},
expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
expires_delta=expires_delta,
)
datetime_expires_at = datetime.datetime.fromtimestamp(
expires_at, datetime.timezone.utc
)
# Set the cookie token
response.set_cookie(
key="token",
value=token,
expires=datetime_expires_at,
httponly=True, # Ensures the cookie is not accessible via JavaScript
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
secure=WEBUI_SESSION_COOKIE_SECURE,
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
secure=WEBUI_SESSION_COOKIE_SECURE,
)
return {
"token": token,
"token_type": "Bearer",
"expires_at": expires_at,
"id": user.id,
"email": user.email,
"name": user.name,
@@ -193,7 +226,7 @@ async def signin(request: Request, response: Response, form_data: SigninForm):
############################
@router.post("/signup", response_model=SigninResponse)
@router.post("/signup", response_model=SessionUserResponse)
async def signup(request: Request, response: Response, form_data: SignupForm):
if WEBUI_AUTH:
if (
@@ -233,18 +266,28 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
)
if user:
expires_delta = parse_duration(request.app.state.config.JWT_EXPIRES_IN)
expires_at = None
if expires_delta:
expires_at = int(time.time()) + int(expires_delta.total_seconds())
token = create_token(
data={"id": user.id},
expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
expires_delta=expires_delta,
)
datetime_expires_at = datetime.datetime.fromtimestamp(
expires_at, datetime.timezone.utc
)
# Set the cookie token
response.set_cookie(
key="token",
value=token,
expires=datetime_expires_at,
httponly=True, # Ensures the cookie is not accessible via JavaScript
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
secure=WEBUI_SESSION_COOKIE_SECURE,
samesite=WEBUI_SESSION_COOKIE_SAME_SITE,
secure=WEBUI_SESSION_COOKIE_SECURE,
)
if request.app.state.config.WEBHOOK_URL:
@@ -261,6 +304,7 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
return {
"token": token,
"token_type": "Bearer",
"expires_at": expires_at,
"id": user.id,
"email": user.email,
"name": user.name,