Spaces:

nathycodes
/

presentercoach

Paused

File size: 9,091 Bytes

ec17af3

import hashlib
import uuid
from datetime import datetime, timedelta

import bcrypt
from fastapi import Body, Cookie, Depends, Header, HTTPException, Query, Response, status
from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
from pydantic import BaseModel
from sqlalchemy.orm import Session

from app import models
from app.config import settings
from app.db import get_db

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/auth/login", auto_error=False)
REFRESH_COOKIE_NAME = "refresh_token"
ACCESS_HEADER_NAME = "X-Presenter-Token"
REFRESH_HEADER_NAME = "X-Refresh-Token"


class TokenPayload(BaseModel):
    sub: str
    email: str | None = None
    role: str = "user"
    jti: str
    type: str


def hash_password(password: str) -> str:
    return bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt(rounds=12)).decode("utf-8")


def verify_password(password: str, hashed_password: str) -> bool:
    return bcrypt.checkpw(password.encode("utf-8"), hashed_password.encode("utf-8"))


def _utcnow() -> datetime:
    return datetime.utcnow()


def _hash_jti(jti: str) -> str:
    return hashlib.sha256(jti.encode("utf-8")).hexdigest()


def _token_expire_delta(token_type: str) -> timedelta:
    if token_type == "refresh":
        return timedelta(days=settings.refresh_token_expire_days)
    return timedelta(minutes=settings.access_token_expire_minutes)


def _create_token(user: models.User, token_type: str, jti: str | None = None) -> tuple[str, str, datetime]:
    token_jti = jti or str(uuid.uuid4())
    expires_at = _utcnow() + _token_expire_delta(token_type)
    payload = {
        "sub": str(user.id),
        "email": user.email,
        "role": "admin" if user.is_admin else "user",
        "iat": _utcnow(),
        "exp": expires_at,
        "jti": token_jti,
        "type": token_type,
    }
    token = jwt.encode(payload, settings.jwt_secret, algorithm=settings.jwt_algorithm)
    return token, token_jti, expires_at


def create_access_token(user: models.User) -> str:
    token, _jti, _expires_at = _create_token(user, "access")
    return token


def create_refresh_token(db: Session, user: models.User) -> str:
    token, jti, expires_at = _create_token(user, "refresh")
    db.add(
        models.RefreshToken(
            user_id=user.id,
            jti_hash=_hash_jti(jti),
            expires_at=expires_at,
        )
    )
    db.commit()
    return token


def set_refresh_cookie(response: Response, token: str) -> None:
    response.set_cookie(
        key=REFRESH_COOKIE_NAME,
        value=token,
        httponly=True,
        secure=settings.cookie_secure,
        samesite="lax",
        max_age=settings.refresh_token_expire_days * 24 * 60 * 60,
        path=settings.cookie_path,
    )


def clear_refresh_cookie(response: Response) -> None:
    response.delete_cookie(
        key=REFRESH_COOKIE_NAME,
        httponly=True,
        secure=settings.cookie_secure,
        samesite="lax",
        path=settings.cookie_path,
    )


def decode_token(token: str, expected_type: str = "access") -> TokenPayload:
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Invalid or expired token",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(
            token,
            settings.jwt_secret,
            algorithms=[settings.jwt_algorithm],
            options={"require_exp": True, "require_iat": True},
        )
    except JWTError as exc:
        raise credentials_exception from exc

    if payload.get("type") != expected_type:
        raise credentials_exception

    try:
        return TokenPayload(
            sub=str(payload["sub"]),
            email=payload.get("email"),
            role=payload.get("role", "user"),
            jti=str(payload["jti"]),
            type=str(payload["type"]),
        )
    except (KeyError, ValueError) as exc:
        raise credentials_exception from exc


def authenticate_user(db: Session, email: str, password: str) -> models.User | None:
    user = db.query(models.User).filter(models.User.email == email).first()
    if not user or not verify_password(password, user.hashed_password):
        return None
    return user


def issue_token_pair(db: Session, response: Response, user: models.User) -> dict:
    access_token = create_access_token(user)
    refresh_token = create_refresh_token(db, user)
    set_refresh_cookie(response, refresh_token)
    return {
        "access_token": access_token,
        "refresh_token": refresh_token,
        "token_type": "bearer",
        "expires_in": settings.access_token_expire_minutes * 60,
        "user": {"id": user.id, "email": user.email, "is_admin": user.is_admin},
    }


def rotate_refresh_token(
    db: Session,
    response: Response,
    refresh_token: str | None,
) -> dict:
    if not refresh_token:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing refresh token")

    payload = decode_token(refresh_token, expected_type="refresh")
    token_row = (
        db.query(models.RefreshToken)
        .filter(models.RefreshToken.jti_hash == _hash_jti(payload.jti))
        .first()
    )
    if token_row is None or token_row.revoked_at is not None or token_row.expires_at <= _utcnow():
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid refresh token")

    try:
        user_id = int(payload.sub)
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid refresh token") from exc

    user = db.query(models.User).filter(models.User.id == user_id).first()
    if user is None:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid refresh token")

    access_token = create_access_token(user)
    new_refresh_token, new_jti, new_expires_at = _create_token(user, "refresh")
    token_row.revoked_at = _utcnow()
    token_row.replaced_by_jti_hash = _hash_jti(new_jti)
    db.add(
        models.RefreshToken(
            user_id=user.id,
            jti_hash=_hash_jti(new_jti),
            expires_at=new_expires_at,
        )
    )
    db.commit()
    set_refresh_cookie(response, new_refresh_token)
    return {
        "access_token": access_token,
        "refresh_token": new_refresh_token,
        "token_type": "bearer",
        "expires_in": settings.access_token_expire_minutes * 60,
        "user": {"id": user.id, "email": user.email, "is_admin": user.is_admin},
    }


def revoke_refresh_token(db: Session, response: Response, refresh_token: str | None) -> None:
    if refresh_token:
        try:
            payload = decode_token(refresh_token, expected_type="refresh")
        except HTTPException:
            payload = None
        if payload is not None:
            token_row = (
                db.query(models.RefreshToken)
                .filter(models.RefreshToken.jti_hash == _hash_jti(payload.jti))
                .first()
            )
            if token_row is not None and token_row.revoked_at is None:
                token_row.revoked_at = _utcnow()
                db.commit()
    clear_refresh_cookie(response)


def get_current_user(
    db: Session = Depends(get_db),
    bearer_token: str | None = Depends(oauth2_scheme),
    header_token: str | None = Header(default=None, alias=ACCESS_HEADER_NAME),
    query_token: str | None = Query(default=None, alias="_pt"),
) -> models.User:
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    candidates = [tok for tok in (header_token, query_token, bearer_token) if tok]
    if not candidates:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Not authenticated",
            headers={"WWW-Authenticate": "Bearer"},
        )

    payload: TokenPayload | None = None
    last_error: HTTPException = credentials_exception
    for candidate in candidates:
        try:
            payload = decode_token(candidate, expected_type="access")
            break
        except HTTPException as exc:
            last_error = exc

    if payload is None:
        raise last_error

    try:
        user_id = int(payload.sub)
    except ValueError as exc:
        raise credentials_exception from exc

    user = db.query(models.User).filter(models.User.id == user_id).first()
    if user is None:
        raise credentials_exception
    return user


def get_refresh_cookie(
    refresh_token: str | None = Cookie(default=None, alias=REFRESH_COOKIE_NAME),
    header_refresh_token: str | None = Header(default=None, alias=REFRESH_HEADER_NAME),
    body_payload: dict | None = Body(default=None),
) -> str | None:
    body_token = None
    if isinstance(body_payload, dict):
        raw = body_payload.get("refresh_token")
        if isinstance(raw, str) and raw.strip():
            body_token = raw.strip()
    return refresh_token or header_refresh_token or body_token