mirror of
https://github.com/open-webui/open-webui
synced 2025-06-26 18:26:48 +00:00
refac
This commit is contained in:
@@ -1,6 +1,9 @@
|
||||
import logging
|
||||
import uuid
|
||||
import jwt
|
||||
import base64
|
||||
import hmac
|
||||
import hashlib
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Optional, Union, List, Dict
|
||||
@@ -8,7 +11,7 @@ from typing import Optional, Union, List, Dict
|
||||
from open_webui.models.users import Users
|
||||
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.env import WEBUI_SECRET_KEY
|
||||
from open_webui.env import WEBUI_SECRET_KEY, TRUSTED_SIGNATURE_KEY
|
||||
|
||||
from fastapi import Depends, HTTPException, Request, Response, status
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
@@ -24,6 +27,23 @@ ALGORITHM = "HS256"
|
||||
# Auth Utils
|
||||
##############
|
||||
|
||||
|
||||
def verify_signature(payload: str, signature: str) -> bool:
|
||||
"""
|
||||
Verifies the HMAC signature of the received payload.
|
||||
"""
|
||||
try:
|
||||
expected_signature = base64.b64encode(
|
||||
hmac.new(TRUSTED_SIGNATURE_KEY, payload.encode(), hashlib.sha256).digest()
|
||||
).decode()
|
||||
|
||||
# Compare securely to prevent timing attacks
|
||||
return hmac.compare_digest(expected_signature, signature)
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
bearer_security = HTTPBearer(auto_error=False)
|
||||
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
||||
|
||||
|
||||
@@ -1008,6 +1008,7 @@ async def process_chat_response(
|
||||
webhook_url = Users.get_user_webhook_url_by_id(user.id)
|
||||
if webhook_url:
|
||||
post_webhook(
|
||||
request.app.state.WEBUI_NAME,
|
||||
webhook_url,
|
||||
f"{title} - {request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}\n\n{content}",
|
||||
{
|
||||
@@ -1873,6 +1874,7 @@ async def process_chat_response(
|
||||
webhook_url = Users.get_user_webhook_url_by_id(user.id)
|
||||
if webhook_url:
|
||||
post_webhook(
|
||||
request.app.state.WEBUI_NAME,
|
||||
webhook_url,
|
||||
f"{title} - {request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}\n\n{content}",
|
||||
{
|
||||
|
||||
@@ -36,7 +36,11 @@ from open_webui.config import (
|
||||
AppConfig,
|
||||
)
|
||||
from open_webui.constants import ERROR_MESSAGES, WEBHOOK_MESSAGES
|
||||
from open_webui.env import WEBUI_AUTH_COOKIE_SAME_SITE, WEBUI_AUTH_COOKIE_SECURE
|
||||
from open_webui.env import (
|
||||
WEBUI_NAME,
|
||||
WEBUI_AUTH_COOKIE_SAME_SITE,
|
||||
WEBUI_AUTH_COOKIE_SECURE,
|
||||
)
|
||||
from open_webui.utils.misc import parse_duration
|
||||
from open_webui.utils.auth import get_password_hash, create_token
|
||||
from open_webui.utils.webhook import post_webhook
|
||||
@@ -66,8 +70,9 @@ auth_manager_config.JWT_EXPIRES_IN = JWT_EXPIRES_IN
|
||||
|
||||
|
||||
class OAuthManager:
|
||||
def __init__(self):
|
||||
def __init__(self, app):
|
||||
self.oauth = OAuth()
|
||||
self.app = app
|
||||
for _, provider_config in OAUTH_PROVIDERS.items():
|
||||
provider_config["register"](self.oauth)
|
||||
|
||||
@@ -200,7 +205,7 @@ class OAuthManager:
|
||||
id=group_model.id, form_data=update_form, overwrite=False
|
||||
)
|
||||
|
||||
async def handle_login(self, provider, request):
|
||||
async def handle_login(self, request, provider):
|
||||
if provider not in OAUTH_PROVIDERS:
|
||||
raise HTTPException(404)
|
||||
# If the provider has a custom redirect URL, use that, otherwise automatically generate one
|
||||
@@ -212,7 +217,7 @@ class OAuthManager:
|
||||
raise HTTPException(404)
|
||||
return await client.authorize_redirect(request, redirect_uri)
|
||||
|
||||
async def handle_callback(self, provider, request, response):
|
||||
async def handle_callback(self, request, provider, response):
|
||||
if provider not in OAUTH_PROVIDERS:
|
||||
raise HTTPException(404)
|
||||
client = self.get_client(provider)
|
||||
@@ -266,6 +271,17 @@ class OAuthManager:
|
||||
Users.update_user_role_by_id(user.id, determined_role)
|
||||
|
||||
if not user:
|
||||
user_count = Users.get_num_users()
|
||||
|
||||
if (
|
||||
request.app.state.USER_COUNT
|
||||
and user_count >= request.app.state.USER_COUNT
|
||||
):
|
||||
raise HTTPException(
|
||||
403,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
# If the user does not exist, check if signups are enabled
|
||||
if auth_manager_config.ENABLE_OAUTH_SIGNUP:
|
||||
# Check if an existing user with the same email already exists
|
||||
@@ -334,6 +350,7 @@ class OAuthManager:
|
||||
|
||||
if auth_manager_config.WEBHOOK_URL:
|
||||
post_webhook(
|
||||
WEBUI_NAME,
|
||||
auth_manager_config.WEBHOOK_URL,
|
||||
WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
|
||||
{
|
||||
@@ -380,6 +397,3 @@ class OAuthManager:
|
||||
# Redirect back to the frontend with the JWT token
|
||||
redirect_url = f"{request.base_url}auth#token={jwt_token}"
|
||||
return RedirectResponse(url=redirect_url, headers=response.headers)
|
||||
|
||||
|
||||
oauth_manager = OAuthManager()
|
||||
|
||||
@@ -2,14 +2,14 @@ import json
|
||||
import logging
|
||||
|
||||
import requests
|
||||
from open_webui.config import WEBUI_FAVICON_URL, WEBUI_NAME
|
||||
from open_webui.config import WEBUI_FAVICON_URL
|
||||
from open_webui.env import SRC_LOG_LEVELS, VERSION
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.setLevel(SRC_LOG_LEVELS["WEBHOOK"])
|
||||
|
||||
|
||||
def post_webhook(url: str, message: str, event_data: dict) -> bool:
|
||||
def post_webhook(name: str, url: str, message: str, event_data: dict) -> bool:
|
||||
try:
|
||||
log.debug(f"post_webhook: {url}, {message}, {event_data}")
|
||||
payload = {}
|
||||
@@ -39,7 +39,7 @@ def post_webhook(url: str, message: str, event_data: dict) -> bool:
|
||||
"sections": [
|
||||
{
|
||||
"activityTitle": message,
|
||||
"activitySubtitle": f"{WEBUI_NAME} ({VERSION}) - {action}",
|
||||
"activitySubtitle": f"{name} ({VERSION}) - {action}",
|
||||
"activityImage": WEBUI_FAVICON_URL,
|
||||
"facts": facts,
|
||||
"markdown": True,
|
||||
|
||||
Reference in New Issue
Block a user