iridescentX/openui

1from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials2from fastapi import HTTPException, status, Depends, Request3from sqlalchemy.orm import Session4 5from apps.webui.models.users import Users6 7from pydantic import BaseModel8from typing import Union, Optional9from constants import ERROR_MESSAGES10from passlib.context import CryptContext11from datetime import datetime, timedelta12import requests13import jwt14import uuid15import logging16import config17 18logging.getLogger("passlib").setLevel(logging.ERROR)19 20 21SESSION_SECRET = config.WEBUI_SECRET_KEY22ALGORITHM = "HS256"23 24##############25# Auth Utils26##############27 28bearer_security = HTTPBearer(auto_error=False)29pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")30 31 32def verify_password(plain_password, hashed_password):33    return (34        pwd_context.verify(plain_password, hashed_password) if hashed_password else None35    )36 37 38def get_password_hash(password):39    return pwd_context.hash(password)40 41 42def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str:43    payload = data.copy()44 45    if expires_delta:46        expire = datetime.utcnow() + expires_delta47        payload.update({"exp": expire})48 49    encoded_jwt = jwt.encode(payload, SESSION_SECRET, algorithm=ALGORITHM)50    return encoded_jwt51 52 53def decode_token(token: str) -> Optional[dict]:54    try:55        decoded = jwt.decode(token, SESSION_SECRET, algorithms=[ALGORITHM])56        return decoded57    except Exception as e:58        return None59 60 61def extract_token_from_auth_header(auth_header: str):62    return auth_header[len("Bearer ") :]63 64 65def create_api_key():66    key = str(uuid.uuid4()).replace("-", "")67    return f"sk-{key}"68 69 70def get_http_authorization_cred(auth_header: str):71    try:72        scheme, credentials = auth_header.split(" ")73        return HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)74    except:75        raise ValueError(ERROR_MESSAGES.INVALID_TOKEN)76 77 78def get_current_user(79    request: Request,80    auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),81):82    token = None83 84    if auth_token is not None:85        token = auth_token.credentials86 87    if token is None and "token" in request.cookies:88        token = request.cookies.get("token")89 90    if token is None:91        raise HTTPException(status_code=403, detail="Not authenticated")92 93    # auth by api key94    if token.startswith("sk-"):95        return get_current_user_by_api_key(token)96 97    # auth by jwt token98    data = decode_token(token)99    if data != None and "id" in data:100        user = Users.get_user_by_id(data["id"])101        if user is None:102            raise HTTPException(103                status_code=status.HTTP_401_UNAUTHORIZED,104                detail=ERROR_MESSAGES.INVALID_TOKEN,105            )106        else:107            Users.update_user_last_active_by_id(user.id)108        return user109    else:110        raise HTTPException(111            status_code=status.HTTP_401_UNAUTHORIZED,112            detail=ERROR_MESSAGES.UNAUTHORIZED,113        )114 115 116def get_current_user_by_api_key(api_key: str):117    user = Users.get_user_by_api_key(api_key)118 119    if user is None:120        raise HTTPException(121            status_code=status.HTTP_401_UNAUTHORIZED,122            detail=ERROR_MESSAGES.INVALID_TOKEN,123        )124    else:125        Users.update_user_last_active_by_id(user.id)126 127    return user128 129 130def get_verified_user(user=Depends(get_current_user)):131    if user.role not in {"user", "admin"}:132        raise HTTPException(133            status_code=status.HTTP_401_UNAUTHORIZED,134            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,135        )136    return user137 138 139def get_admin_user(user=Depends(get_current_user)):140    if user.role != "admin":141        raise HTTPException(142            status_code=status.HTTP_401_UNAUTHORIZED,143            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,144        )145    return user146