utils.py 2.03 KB
Newer Older
1
2
3
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi import HTTPException, status, Depends
from apps.web.models.users import Users
4
5
from pydantic import BaseModel
from typing import Union, Optional
6
from constants import ERROR_MESSAGES
7
8
9
10
from passlib.context import CryptContext
from datetime import datetime, timedelta
import requests
import jwt
Timothy J. Baek's avatar
Timothy J. Baek committed
11
import logging
12
13
import config

Timothy J. Baek's avatar
Timothy J. Baek committed
14
15
16
logging.getLogger("passlib").setLevel(logging.ERROR)


17
SESSION_SECRET = config.WEBUI_SECRET_KEY
18
19
20
21
22
23
24
25
26
27
28
ALGORITHM = "HS256"

##############
# Auth Utils
##############

bearer_scheme = HTTPBearer()
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")


def verify_password(plain_password, hashed_password):
Timothy J. Baek's avatar
Timothy J. Baek committed
29
30
31
    return (
        pwd_context.verify(plain_password, hashed_password) if hashed_password else None
    )
32
33
34
35
36
37


def get_password_hash(password):
    return pwd_context.hash(password)


Timothy J. Baek's avatar
Timothy J. Baek committed
38
def create_token(data: dict, expires_delta: Union[timedelta, None] = None) -> str:
39
40
41
42
43
44
    payload = data.copy()

    if expires_delta:
        expire = datetime.utcnow() + expires_delta
        payload.update({"exp": expire})

45
    encoded_jwt = jwt.encode(payload, SESSION_SECRET, algorithm=ALGORITHM)
46
47
48
49
50
    return encoded_jwt


def decode_token(token: str) -> Optional[dict]:
    try:
51
        decoded = jwt.decode(token, SESSION_SECRET)
52
53
54
55
56
57
        return decoded
    except Exception as e:
        return None


def extract_token_from_auth_header(auth_header: str):
Timothy J. Baek's avatar
Timothy J. Baek committed
58
    return auth_header[len("Bearer ") :]
59
60


Timothy J. Baek's avatar
Timothy J. Baek committed
61
def get_current_user(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
62
    data = decode_token(auth_token.credentials)
63
64
    if data != None and "id" in data:
        user = Users.get_user_by_id(data["id"])
65
66
67
68
        if user is None:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.INVALID_TOKEN,
69
            )
70
        return user
71
72
73
74
75
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.UNAUTHORIZED,
        )