open-webui/backend/utils/utils.py

126 lines
3.3 KiB
Python
Raw Normal View History

from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi import HTTPException, status, Depends
2024-04-02 19:05:53 +02:00
from apps.web.models.users import Users
2024-03-26 11:22:17 +01:00
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-03-26 11:22:17 +01:00
import uuid
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)
SESSION_SECRET = config.WEBUI_SECRET_KEY
2023-11-19 01:47:12 +01:00
ALGORITHM = "HS256"
##############
# Auth Utils
##############
2024-02-01 21:52:11 +01:00
bearer_security = HTTPBearer()
2023-11-19 01:47:12 +01:00
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, SESSION_SECRET, algorithm=ALGORITHM)
2023-11-19 01:47:12 +01:00
return encoded_jwt
def decode_token(token: str) -> Optional[dict]:
try:
decoded = jwt.decode(token, SESSION_SECRET, algorithms=[ALGORITHM])
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-03-26 11:22:17 +01:00
def create_api_key():
key = str(uuid.uuid4()).replace("-", "")
return f"sk-{key}"
2024-02-24 07:44:56 +01:00
def get_http_authorization_cred(auth_header: str):
try:
scheme, credentials = auth_header.split(" ")
2024-02-25 07:10:43 +01:00
return HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)
2024-02-24 07:44:56 +01:00
except:
raise ValueError(ERROR_MESSAGES.INVALID_TOKEN)
2024-02-11 02:54:33 +01:00
def get_current_user(
auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
):
2024-03-26 11:22:17 +01:00
# auth by api key
if auth_token.credentials.startswith("sk-"):
return get_current_user_by_api_key(auth_token.credentials)
# auth by jwt token
data = decode_token(auth_token.credentials)
if data != None and "id" in data:
user = Users.get_user_by_id(data["id"])
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,
)
2024-04-02 18:42:45 +02:00
2024-03-26 11:22:17 +01:00
def get_current_user_by_api_key(api_key: str):
2024-04-02 19:06:58 +02:00
user = Users.get_user_by_api_key(api_key)
2024-03-26 11:22:17 +01:00
if user is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.INVALID_TOKEN,
)
return user
2024-04-02 18:42:45 +02:00
2024-02-11 02:54:33 +01:00
def get_verified_user(user=Depends(get_current_user)):
if user.role not in {"user", "admin"}:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
)
2024-02-11 02:54:33 +01:00
return user
2024-02-11 02:54:33 +01:00
def get_admin_user(user=Depends(get_current_user)):
if user.role != "admin":
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
)
2024-02-11 02:54:33 +01:00
return user