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