open-webui/backend/utils/utils.py

76 lines
2.1 KiB
Python
Raw Normal View History

from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi import HTTPException, status, Depends
from apps.web.models.users import Users
2023-11-19 01:47:12 +01:00
from pydantic import BaseModel
from typing import Union, Optional
from constants import ERROR_MESSAGES
2023-11-19 01:47:12 +01:00
from passlib.context import CryptContext
from datetime import datetime, timedelta
import requests
import jwt
2024-01-05 21:22:27 +01:00
import logging
2023-11-19 01:47:12 +01:00
import config
2024-01-05 21:22:27 +01:00
logging.getLogger("passlib").setLevel(logging.ERROR)
JWT_SECRET_KEY = config.WEBUI_JWT_SECRET_KEY
2023-11-19 01:47:12 +01:00
ALGORITHM = "HS256"
##############
# Auth Utils
##############
bearer_scheme = HTTPBearer()
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
def verify_password(plain_password, hashed_password):
2024-01-05 21:22:27 +01:00
return (
pwd_context.verify(plain_password, hashed_password) if hashed_password else None
)
2023-11-19 01:47:12 +01:00
def get_password_hash(password):
return pwd_context.hash(password)
2024-01-05 21:22:27 +01:00
def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str:
2023-11-19 01:47:12 +01:00
payload = data.copy()
if expires_delta:
expire = datetime.utcnow() + expires_delta
payload.update({"exp": expire})
encoded_jwt = jwt.encode(payload, JWT_SECRET_KEY, algorithm=ALGORITHM)
return encoded_jwt
def decode_token(token: str) -> Optional[dict]:
try:
2024-01-05 21:22:27 +01:00
decoded = jwt.decode(token, JWT_SECRET_KEY, options={"verify_signature": False})
2023-11-19 01:47:12 +01:00
return decoded
except Exception as e:
return None
def extract_token_from_auth_header(auth_header: str):
2024-01-05 21:22:27 +01:00
return auth_header[len("Bearer ") :]
2023-11-19 01:47:12 +01:00
2024-01-05 21:22:27 +01:00
def get_current_user(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
data = decode_token(auth_token.credentials)
if data != None and "email" in data:
user = Users.get_user_by_email(data["email"])
if user is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.INVALID_TOKEN,
2023-11-19 01:47:12 +01:00
)
return user
else:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.UNAUTHORIZED,
)