gitscope
https://dev.kronoz.win/r/marzneshin.git
app/dependencies.py · v0.3.0-alpha9
Source
1 41ba8cb Dawsh 2y ago from datetime import datetime
2 c13a1aa Dawsh 2y ago from typing import Annotated
3 c13a1aa Dawsh 2y ago
4 c13a1aa Dawsh 2y ago from fastapi import Depends, HTTPException, status
5 f7af2b3 Dawsh 2y ago from sqlalchemy.exc import SQLAlchemyError
6 c13a1aa Dawsh 2y ago from sqlalchemy.orm import Session
7 c13a1aa Dawsh 2y ago
8 f7af2b3 Dawsh 2y ago from app.db import crud, User, GetDB
9 c13a1aa Dawsh 2y ago from app.models.admin import Admin, oauth2_scheme
10 c13a1aa Dawsh 2y ago from app.utils.jwt import get_admin_payload
11 c13a1aa Dawsh 2y ago from config import SUDOERS
12 c13a1aa Dawsh 2y ago
13 c13a1aa Dawsh 2y ago
14 f7af2b3 Dawsh 2y ago def get_db(): # Dependency
15 f7af2b3 Dawsh 2y ago with GetDB() as db:
16 f7af2b3 Dawsh 2y ago yield db
17 f7af2b3 Dawsh 2y ago
18 f7af2b3 Dawsh 2y ago
19 01cdf2e Dawsh 2y ago def get_admin(
20 01cdf2e Dawsh 2y ago db: Annotated[Session, Depends(get_db)],
21 01cdf2e Dawsh 2y ago token: Annotated[str, Depends(oauth2_scheme)],
22 01cdf2e Dawsh 2y ago ):
23 c13a1aa Dawsh 2y ago payload = get_admin_payload(token)
24 c13a1aa Dawsh 2y ago if not payload:
25 c13a1aa Dawsh 2y ago return
26 c13a1aa Dawsh 2y ago
27 01cdf2e Dawsh 2y ago if payload["username"] in SUDOERS and payload["is_sudo"] is True:
28 01cdf2e Dawsh 2y ago return Admin(username=payload["username"], is_sudo=True)
29 c13a1aa Dawsh 2y ago
30 01cdf2e Dawsh 2y ago dbadmin = crud.get_admin(db, payload["username"])
31 c13a1aa Dawsh 2y ago if not dbadmin:
32 c13a1aa Dawsh 2y ago return
33 c13a1aa Dawsh 2y ago
34 c13a1aa Dawsh 2y ago if dbadmin.password_reset_at:
35 c13a1aa Dawsh 2y ago if not payload.get("created_at"):
36 c13a1aa Dawsh 2y ago return
37 c13a1aa Dawsh 2y ago if dbadmin.password_reset_at > payload.get("created_at"):
38 c13a1aa Dawsh 2y ago return
39 c13a1aa Dawsh 2y ago
40 c13a1aa Dawsh 2y ago return Admin.model_validate(dbadmin)
41 c13a1aa Dawsh 2y ago
42 c13a1aa Dawsh 2y ago
43 c13a1aa Dawsh 2y ago def get_current_admin(admin: Annotated[Admin, Depends(get_admin)]):
44 c13a1aa Dawsh 2y ago if not admin:
45 c13a1aa Dawsh 2y ago raise HTTPException(
46 c13a1aa Dawsh 2y ago status_code=status.HTTP_401_UNAUTHORIZED,
47 c13a1aa Dawsh 2y ago detail="Could not validate credentials",
48 c13a1aa Dawsh 2y ago headers={"WWW-Authenticate": "Bearer"},
49 c13a1aa Dawsh 2y ago )
50 c13a1aa Dawsh 2y ago
51 c13a1aa Dawsh 2y ago return admin
52 c13a1aa Dawsh 2y ago
53 c13a1aa Dawsh 2y ago
54 c13a1aa Dawsh 2y ago def sudo_admin(admin: Annotated[Admin, Depends(get_current_admin)]):
55 c13a1aa Dawsh 2y ago if not admin.is_sudo:
56 c13a1aa Dawsh 2y ago raise HTTPException(
57 c13a1aa Dawsh 2y ago status_code=status.HTTP_403_FORBIDDEN,
58 c13a1aa Dawsh 2y ago detail="Access Denied",
59 c13a1aa Dawsh 2y ago )
60 f7af2b3 Dawsh 2y ago return admin
61 c13a1aa Dawsh 2y ago
62 c13a1aa Dawsh 2y ago
63 01cdf2e Dawsh 2y ago def get_subscription_user(
64 01cdf2e Dawsh 2y ago username: str, key: str, db: Annotated[Session, Depends(get_db)]
65 01cdf2e Dawsh 2y ago ):
66 c13a1aa Dawsh 2y ago try:
67 c13a1aa Dawsh 2y ago int(key, 16)
68 c13a1aa Dawsh 2y ago except ValueError:
69 c13a1aa Dawsh 2y ago raise HTTPException(404)
70 c13a1aa Dawsh 2y ago
71 c13a1aa Dawsh 2y ago db_user = crud.get_user(db, username)
72 c13a1aa Dawsh 2y ago if not db_user or db_user.key != key:
73 c13a1aa Dawsh 2y ago raise HTTPException(404)
74 c13a1aa Dawsh 2y ago return db_user
75 c13a1aa Dawsh 2y ago
76 c13a1aa Dawsh 2y ago
77 01cdf2e Dawsh 2y ago def get_user(
78 01cdf2e Dawsh 2y ago username: str,
79 01cdf2e Dawsh 2y ago admin: Annotated[Admin, Depends(get_current_admin)],
80 01cdf2e Dawsh 2y ago db: Annotated[Session, Depends(get_db)],
81 01cdf2e Dawsh 2y ago ):
82 f7af2b3 Dawsh 2y ago db_user = crud.get_user(db, username)
83 01cdf2e Dawsh 2y ago if not (
84 01cdf2e Dawsh 2y ago admin.is_sudo or (db_user.admin and db_user.admin.username == admin.username)
85 01cdf2e Dawsh 2y ago ):
86 f7af2b3 Dawsh 2y ago raise HTTPException(status_code=403, detail="You're not allowed")
87 f7af2b3 Dawsh 2y ago
88 f7af2b3 Dawsh 2y ago if not db_user:
89 f7af2b3 Dawsh 2y ago raise HTTPException(status_code=404, detail="User not found")
90 f7af2b3 Dawsh 2y ago
91 f7af2b3 Dawsh 2y ago return db_user
92 f7af2b3 Dawsh 2y ago
93 f7af2b3 Dawsh 2y ago
94 41ba8cb Dawsh 2y ago def parse_start_date(start: str | None = None):
95 41ba8cb Dawsh 2y ago if not start:
96 41ba8cb Dawsh 2y ago return datetime.fromtimestamp(datetime.utcnow().timestamp() - 30 * 24 * 3600)
97 41ba8cb Dawsh 2y ago else:
98 41ba8cb Dawsh 2y ago return datetime.fromisoformat(start)
99 41ba8cb Dawsh 2y ago
100 41ba8cb Dawsh 2y ago
101 41ba8cb Dawsh 2y ago def parse_end_date(end: str | None = None):
102 41ba8cb Dawsh 2y ago if not end:
103 41ba8cb Dawsh 2y ago return datetime.utcnow()
104 41ba8cb Dawsh 2y ago else:
105 41ba8cb Dawsh 2y ago return datetime.fromisoformat(end)
106 41ba8cb Dawsh 2y ago
107 41ba8cb Dawsh 2y ago
108 f7af2b3 Dawsh 2y ago SubUserDep = Annotated[User, Depends(get_subscription_user)]
109 f7af2b3 Dawsh 2y ago UserDep = Annotated[User, Depends(get_user)]
110 c13a1aa Dawsh 2y ago AdminDep = Annotated[Admin, Depends(get_current_admin)]
111 c13a1aa Dawsh 2y ago SudoAdminDep = Annotated[Admin, Depends(sudo_admin)]
112 c13a1aa Dawsh 2y ago DBDep = Annotated[Session, Depends(get_db)]
113 41ba8cb Dawsh 2y ago StartDateDep = Annotated[datetime, Depends(parse_start_date)]
114 41ba8cb Dawsh 2y ago EndDateDep = Annotated[datetime, Depends(parse_end_date)]