gitscope
← marzneshin / v0.1.0 / 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 = {"sub": username, "access": "sudo" if is_sudo else "admin", "iat": datetime.utcnow()}
19    if JWT_ACCESS_TOKEN_EXPIRE_MINUTES > 0:
20        expire = datetime.utcnow() + timedelta(minutes=JWT_ACCESS_TOKEN_EXPIRE_MINUTES)
21        data["exp"] = expire
22    encoded_jwt = jwt.encode(data, get_secret_key(), algorithm="HS256")
23    return encoded_jwt
24
25
26def get_admin_payload(token: str) -> Union[dict, None]:
27    try:
28        payload = jwt.decode(token, get_secret_key(), algorithms=["HS256"])
29        username: str = payload.get("sub")
30        access: str = payload.get("access")
31        if not username or access not in ('admin', 'sudo'):
32            return
33        try:
34            created_at = datetime.utcfromtimestamp(payload['iat'])
35        except KeyError:
36            created_at = None
37
38        return {"username": username, "is_sudo": access == "sudo", "created_at": created_at}
39    except JWTError:
40        return