utils.py 2.82 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
ALGORITHM = "HS256"

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

Tim Farrell's avatar
Tim Farrell committed
24
bearer_security = HTTPBearer()
25
26
27
28
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, algorithms=[ALGORITHM])
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
62
63
def get_http_authorization_cred(auth_header: str):
    try:
        scheme, credentials = auth_header.split(" ")
Timothy J. Baek's avatar
Timothy J. Baek committed
64
        return HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)
Timothy J. Baek's avatar
Timothy J. Baek committed
65
66
67
68
    except:
        raise ValueError(ERROR_MESSAGES.INVALID_TOKEN)


Timothy J. Baek's avatar
Timothy J. Baek committed
69
70
71
def get_current_user(
    auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
):
72
    data = decode_token(auth_token.credentials)
73
74
    if data != None and "id" in data:
        user = Users.get_user_by_id(data["id"])
75
76
77
78
        if user is None:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.INVALID_TOKEN,
79
            )
80
        return user
81
82
83
84
85
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.UNAUTHORIZED,
        )
86
87


Timothy J. Baek's avatar
Timothy J. Baek committed
88
def get_verified_user(user=Depends(get_current_user)):
89
90
91
92
93
    if user.role not in {"user", "admin"}:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
94
    return user
95
96


Timothy J. Baek's avatar
Timothy J. Baek committed
97
def get_admin_user(user=Depends(get_current_user)):
98
99
100
101
102
    if user.role != "admin":
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
103
    return user