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