pdf-quiz-generator/backend/app/routers/auth.py

import secrets
from collections import defaultdict
from datetime import datetime, timedelta

from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks, status, Request
from sqlalchemy.orm import Session

from app.database import get_db
from app.models.user import User
from app.models.email_verification import EmailVerification
from app.models.password_reset import PasswordReset
from app.schemas.auth import (
    UserCreate, UserResponse, Token, LoginRequest,
    UserUpdateMe, ForgotPasswordRequest, ResetPasswordRequest,
)
from app.services import email_service
from app.utils.auth import (
    get_password_hash, verify_password, create_access_token, get_current_user,
)

router = APIRouter()

# Simple in-memory login rate limiter: max 10 attempts per IP per 15 min
_login_attempts: dict = defaultdict(list)

def _check_login_rate_limit(client_ip: str):
    now = datetime.utcnow()
    window = now - timedelta(minutes=15)
    attempts = [t for t in _login_attempts[client_ip] if t > window]
    _login_attempts[client_ip] = attempts
    if len(attempts) >= 10:
        raise HTTPException(status_code=429, detail="Too many login attempts. Try again in 15 minutes.")
    _login_attempts[client_ip].append(now)


# Rate limit: max 3 reset requests per email per hour
RESET_LIMIT = 3
RESET_WINDOW_HOURS = 1


def _check_reset_rate_limit(db: Session, email: str):
    user = db.query(User).filter(User.email == email).first()
    if not user:
        return  # silently ignore unknown emails
    window_start = datetime.utcnow() - timedelta(hours=RESET_WINDOW_HOURS)
    count = db.query(PasswordReset).filter(
        PasswordReset.user_id == user.id,
        PasswordReset.created_at >= window_start,
    ).count()
    if count >= RESET_LIMIT:
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail=f"Too many reset requests. Please wait before trying again.",
        )


@router.post("/register")
async def register(user_data: UserCreate, background_tasks: BackgroundTasks, db: Session = Depends(get_db)):
    if len(user_data.password) < 8:
        raise HTTPException(status_code=400, detail="Password must be at least 8 characters")

    existing = db.query(User).filter(User.email == user_data.email).first()
    if existing:
        raise HTTPException(status_code=400, detail="Email already registered")

    is_first_user = db.query(User).count() == 0
    user = User(
        email=user_data.email,
        hashed_password=get_password_hash(user_data.password),
        name=user_data.name,
        role="admin" if is_first_user else "user",
    )
    db.add(user)
    db.flush()

    # Create email verification record
    token = secrets.token_urlsafe(32)
    verification = EmailVerification(
        user_id=user.id,
        token=token,
        expires_at=datetime.utcnow() + timedelta(hours=24),
        verified_at=datetime.utcnow() if is_first_user else None,
    )
    db.add(verification)
    db.commit()
    db.refresh(user)

    if is_first_user:
        # Auto-verified admin — log them in immediately
        access_token = create_access_token(data={"sub": user.email})
        return {"access_token": access_token, "token_type": "bearer"}

    # Regular user — must verify email before logging in
    background_tasks.add_task(email_service.send_verification_email, user.email, user.name, token)
    return {"requires_verification": True, "message": "Account created. Please check your email to verify your account before logging in."}


@router.post("/login", response_model=Token)
def login(login_data: LoginRequest, db: Session = Depends(get_db), request: Request = None):
    if request:
        client_ip = request.client.host if request.client else "unknown"
        _check_login_rate_limit(client_ip)

    user = db.query(User).filter(User.email == login_data.email).first()
    if not user or not verify_password(login_data.password, user.hashed_password):
        raise HTTPException(status_code=401, detail="Invalid email or password")

    # Check email verification — skip for users without any verification record (legacy/seeded)
    verification = db.query(EmailVerification).filter(EmailVerification.user_id == user.id).first()
    if verification and verification.verified_at is None:
        raise HTTPException(
            status_code=403,
            detail="Email not verified. Please check your inbox and verify your email before logging in.",
        )

    access_token = create_access_token(data={"sub": user.email})
    return Token(access_token=access_token)


@router.get("/verify-email")
def verify_email(token: str, db: Session = Depends(get_db)):
    record = db.query(EmailVerification).filter(EmailVerification.token == token).first()
    if not record:
        raise HTTPException(status_code=400, detail="Invalid verification token")
    if record.verified_at is not None:
        return {"message": "Email already verified. You can log in."}
    if datetime.utcnow() > record.expires_at:
        raise HTTPException(status_code=400, detail="Verification link has expired. Please register again or request a new link.")

    record.verified_at = datetime.utcnow()
    db.commit()
    return {"message": "Email verified successfully! You can now log in."}


@router.post("/resend-verification")
async def resend_verification(data: ForgotPasswordRequest, background_tasks: BackgroundTasks, db: Session = Depends(get_db)):
    user = db.query(User).filter(User.email == data.email).first()
    if not user:
        return {"message": "If that email exists, a verification link has been sent."}

    record = db.query(EmailVerification).filter(EmailVerification.user_id == user.id).first()
    if record and record.verified_at is not None:
        return {"message": "Email already verified."}

    token = secrets.token_urlsafe(32)
    if record:
        record.token = token
        record.expires_at = datetime.utcnow() + timedelta(hours=24)
    else:
        db.add(EmailVerification(user_id=user.id, token=token, expires_at=datetime.utcnow() + timedelta(hours=24)))
    db.commit()

    background_tasks.add_task(email_service.send_verification_email, user.email, user.name, token)
    return {"message": "If that email exists, a verification link has been sent."}


@router.post("/forgot-password")
async def forgot_password(data: ForgotPasswordRequest, background_tasks: BackgroundTasks, db: Session = Depends(get_db)):
    _check_reset_rate_limit(db, data.email)

    user = db.query(User).filter(User.email == data.email).first()
    # Always return same message to avoid email enumeration
    if not user:
        return {"message": "If that email is registered, a reset link has been sent."}

    token = secrets.token_urlsafe(32)
    reset = PasswordReset(
        user_id=user.id,
        token=token,
        expires_at=datetime.utcnow() + timedelta(hours=1),
    )
    db.add(reset)
    db.commit()

    background_tasks.add_task(email_service.send_password_reset_email, user.email, user.name, token)
    return {"message": "If that email is registered, a reset link has been sent."}


@router.post("/reset-password")
def reset_password(data: ResetPasswordRequest, db: Session = Depends(get_db)):
    if len(data.new_password) < 8:
        raise HTTPException(status_code=400, detail="Password must be at least 8 characters")

    record = db.query(PasswordReset).filter(
        PasswordReset.token == data.token,
        PasswordReset.used == False,
    ).first()
    if not record:
        raise HTTPException(status_code=400, detail="Invalid or already used reset token")
    if datetime.utcnow() > record.expires_at:
        raise HTTPException(status_code=400, detail="Reset link has expired. Please request a new one.")

    user = db.query(User).filter(User.id == record.user_id).first()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")

    user.hashed_password = get_password_hash(data.new_password)
    record.used = True
    db.commit()
    return {"message": "Password reset successfully. You can now log in."}


@router.get("/me", response_model=UserResponse)
def get_me(current_user: User = Depends(get_current_user)):
    return current_user


@router.put("/me", response_model=UserResponse)
def update_me(data: UserUpdateMe, db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    if data.new_password:
        if not data.current_password:
            raise HTTPException(status_code=400, detail="Current password required to set a new one")
        if not verify_password(data.current_password, current_user.hashed_password):
            raise HTTPException(status_code=400, detail="Current password is incorrect")
        if len(data.new_password) < 8:
            raise HTTPException(status_code=400, detail="New password must be at least 8 characters")
        current_user.hashed_password = get_password_hash(data.new_password)

    if data.name:
        current_user.name = data.name

    db.commit()
    db.refresh(current_user)
    return current_user