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.3.0-alpha1
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 01cdf2e Dawsh 2y ago data = {
19 01cdf2e Dawsh 2y ago "sub": username,
20 01cdf2e Dawsh 2y ago "access": "sudo" if is_sudo else "admin",
21 01cdf2e Dawsh 2y ago "iat": datetime.utcnow(),
22 01cdf2e Dawsh 2y ago }
23 a47eb4c Dawsh 2y ago if JWT_ACCESS_TOKEN_EXPIRE_MINUTES > 0:
24 a47eb4c Dawsh 2y ago expire = datetime.utcnow() + timedelta(minutes=JWT_ACCESS_TOKEN_EXPIRE_MINUTES)
25 a47eb4c Dawsh 2y ago data["exp"] = expire
26 a47eb4c Dawsh 2y ago encoded_jwt = jwt.encode(data, get_secret_key(), algorithm="HS256")
27 a47eb4c Dawsh 2y ago return encoded_jwt
28 a47eb4c Dawsh 2y ago
29 a47eb4c Dawsh 2y ago
30 a47eb4c Dawsh 2y ago def get_admin_payload(token: str) -> Union[dict, None]:
31 a47eb4c Dawsh 2y ago try:
32 a47eb4c Dawsh 2y ago payload = jwt.decode(token, get_secret_key(), algorithms=["HS256"])
33 a47eb4c Dawsh 2y ago username: str = payload.get("sub")
34 a47eb4c Dawsh 2y ago access: str = payload.get("access")
35 01cdf2e Dawsh 2y ago if not username or access not in ("admin", "sudo"):
36 a47eb4c Dawsh 2y ago return
37 a47eb4c Dawsh 2y ago try:
38 01cdf2e Dawsh 2y ago created_at = datetime.utcfromtimestamp(payload["iat"])
39 a47eb4c Dawsh 2y ago except KeyError:
40 a47eb4c Dawsh 2y ago created_at = None
41 a47eb4c Dawsh 2y ago
42 01cdf2e Dawsh 2y ago return {
43 01cdf2e Dawsh 2y ago "username": username,
44 01cdf2e Dawsh 2y ago "is_sudo": access == "sudo",
45 01cdf2e Dawsh 2y ago "created_at": created_at,
46 01cdf2e Dawsh 2y ago }
47 a47eb4c Dawsh 2y ago except JWTError:
48 a47eb4c Dawsh 2y ago return