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.auth import get_admin_payload
 10
 11
 12def get_db():
 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        created_at = payload.get("created_at")
 31        if not created_at or dbadmin.password_reset_at > created_at:
 32            return
 33
 34    return Admin.model_validate(dbadmin)
 35
 36
 37def get_current_admin(admin: Annotated[Admin, Depends(get_admin)]):
 38    if not admin:
 39        raise HTTPException(
 40            status_code=status.HTTP_401_UNAUTHORIZED,
 41            detail="Could not validate credentials",
 42            headers={"WWW-Authenticate": "Bearer"},
 43        )
 44
 45    return admin
 46
 47
 48def sudo_admin(admin: Annotated[Admin, Depends(get_current_admin)]):
 49    if not admin.is_sudo:
 50        raise HTTPException(
 51            status_code=status.HTTP_403_FORBIDDEN,
 52            detail="Access Denied",
 53        )
 54    return admin
 55
 56
 57def get_subscription_user(
 58    username: str, key: str, db: Annotated[Session, Depends(get_db)]
 59):
 60    try:
 61        int(key, 16)
 62    except ValueError:
 63        raise HTTPException(status_code=404)
 64
 65    db_user = crud.get_user(db, username)
 66    if db_user and db_user.key == key:
 67        return db_user
 68    else:
 69        raise HTTPException(status_code=404)
 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 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(
 92            datetime.utcnow().timestamp() - 30 * 24 * 3600
 93        )
 94    else:
 95        return datetime.fromisoformat(start)
 96
 97
 98def parse_end_date(end: str | None = None):
 99    if not end:
100        return datetime.utcnow()
101    else:
102        return datetime.fromisoformat(end)
103
104
105SubUserDep = Annotated[User, Depends(get_subscription_user)]
106UserDep = Annotated[User, Depends(get_user)]
107AdminDep = Annotated[Admin, Depends(get_current_admin)]
108SudoAdminDep = Annotated[Admin, Depends(sudo_admin)]
109DBDep = Annotated[Session, Depends(get_db)]
110StartDateDep = Annotated[datetime, Depends(parse_start_date)]
111EndDateDep = Annotated[datetime, Depends(parse_end_date)]