gitscope
← marzneshin / v0.3.0-alpha6 / app/utils/jwt.py
https://dev.kronoz.win/r/marzneshin.git
app/utils/jwt.py · Python
Blame Raw
 1from datetime import datetime, timedelta
 2from functools import lru_cache
 3from typing import Union
 4
 5from jose import JWTError, jwt
 6
 7from config import JWT_ACCESS_TOKEN_EXPIRE_MINUTES, SUDOERS
 8from app.db import GetDB, get_jwt_secret_key
 9
10
11@lru_cache(maxsize=None)
12def get_secret_key():
13    with GetDB() as db:
14        return get_jwt_secret_key(db)
15
16
17def create_admin_token(username: str, is_sudo=False) -> str:
18    data = {
19        "sub": username,
20        "access": "sudo" if is_sudo else "admin",
21        "iat": datetime.utcnow(),
22    }
23    if JWT_ACCESS_TOKEN_EXPIRE_MINUTES > 0:
24        expire = datetime.utcnow() + timedelta(minutes=JWT_ACCESS_TOKEN_EXPIRE_MINUTES)
25        data["exp"] = expire
26    encoded_jwt = jwt.encode(data, get_secret_key(), algorithm="HS256")
27    return encoded_jwt
28
29
30def get_admin_payload(token: str) -> Union[dict, None]:
31    try:
32        payload = jwt.decode(token, get_secret_key(), algorithms=["HS256"])
33        username: str = payload.get("sub")
34        access: str = payload.get("access")
35        if not username or access not in ("admin", "sudo"):
36            return
37        try:
38            created_at = datetime.utcfromtimestamp(payload["iat"])
39        except KeyError:
40            created_at = None
41
42        return {
43            "username": username,
44            "is_sudo": access == "sudo",
45            "created_at": created_at,
46        }
47    except JWTError:
48        return