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(db: Annotated[Session, Depends(get_db)],
 20              token: Annotated[str, Depends(oauth2_scheme)]):
 21    payload = get_admin_payload(token)
 22    if not payload:
 23        return
 24
 25    if payload['username'] in SUDOERS and payload['is_sudo'] is True:
 26        return Admin(username=payload['username'], is_sudo=True)
 27
 28    dbadmin = crud.get_admin(db, payload['username'])
 29    if not dbadmin:
 30        return
 31
 32    if dbadmin.password_reset_at:
 33        if not payload.get("created_at"):
 34            return
 35        if dbadmin.password_reset_at > payload.get("created_at"):
 36            return
 37
 38    return Admin.model_validate(dbadmin)
 39
 40
 41def get_current_admin(admin: Annotated[Admin, Depends(get_admin)]):
 42    if not admin:
 43        raise HTTPException(
 44            status_code=status.HTTP_401_UNAUTHORIZED,
 45            detail="Could not validate credentials",
 46            headers={"WWW-Authenticate": "Bearer"},
 47        )
 48
 49    return admin
 50
 51
 52def sudo_admin(admin: Annotated[Admin, Depends(get_current_admin)]):
 53    if not admin.is_sudo:
 54        raise HTTPException(
 55            status_code=status.HTTP_403_FORBIDDEN,
 56            detail="Access Denied",
 57        )
 58    return admin
 59
 60
 61def get_subscription_user(username: str, key: str, db: Annotated[Session, Depends(get_db)]):
 62    try:
 63        int(key, 16)
 64    except ValueError:
 65        raise HTTPException(404)
 66
 67    db_user = crud.get_user(db, username)
 68    if not db_user or db_user.key != key:
 69        raise HTTPException(404)
 70    return db_user
 71
 72
 73def get_user(username: str,
 74             admin: Annotated[Admin, Depends(get_current_admin)],
 75             db: Annotated[Session, Depends(get_db)]):
 76    db_user = crud.get_user(db, username)
 77    if not (admin.is_sudo or (db_user.admin and db_user.admin.username == admin.username)):
 78        raise HTTPException(status_code=403, detail="You're not allowed")
 79
 80    if not db_user:
 81        raise HTTPException(status_code=404, detail="User not found")
 82
 83    return db_user
 84
 85
 86def parse_start_date(start: str | None = None):
 87    if not start:
 88        return datetime.fromtimestamp(datetime.utcnow().timestamp() - 30 * 24 * 3600)
 89    else:
 90        return datetime.fromisoformat(start)
 91
 92
 93def parse_end_date(end: str | None = None):
 94    if not end:
 95        return datetime.utcnow()
 96    else:
 97        return datetime.fromisoformat(end)
 98
 99
100SubUserDep = Annotated[User, Depends(get_subscription_user)]
101UserDep = Annotated[User, Depends(get_user)]
102AdminDep = Annotated[Admin, Depends(get_current_admin)]
103SudoAdminDep = Annotated[Admin, Depends(sudo_admin)]
104DBDep = Annotated[Session, Depends(get_db)]
105StartDateDep = Annotated[datetime, Depends(parse_start_date)]
106EndDateDep = Annotated[datetime, Depends(parse_end_date)]