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