kite/backend/api/auth/tokens.py

from __future__ import annotations

import hashlib
import time
import uuid
from typing import Optional

from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from sqlalchemy import select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from db.models import BackupCode, Quota, RefreshToken, User
from deps.auth import get_current_user
from deps.db import get_db
from deps.utils import get_client_ip
from services import auth as auth_service
from services.audit import write_audit_log
from api.auth.shared import (
    limiter,
    RegisterRequest,
    LoginRequest,
    PreferencesUpdate,
    _set_refresh_cookie,
    _user_dict,
)

router = APIRouter()


# ── POST /register ────────────────────────────────────────────────────────────

@router.post("/register", status_code=status.HTTP_201_CREATED)
@limiter.limit("10/minute")
async def register(
    request: Request,
    body: RegisterRequest,
    session: AsyncSession = Depends(get_db),
):
    try:
        auth_service.validate_password_strength(body.password)
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc))

    if await auth_service.check_hibp(body.password):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="This password has appeared in a data breach. Choose a different password.",
        )

    result = await session.execute(
        select(User).where(
            (User.email == str(body.email)) | (User.handle == body.handle)
        )
    )
    if result.scalar_one_or_none() is not None:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail="Email or handle already in use",
        )

    user_id = uuid.uuid4()
    new_user = User(
        id=user_id,
        handle=body.handle,
        email=str(body.email),
        password_hash=auth_service.hash_password(body.password),
        role="user",
        is_active=True,
        password_must_change=False,
    )
    quota = Quota(
        user_id=user_id,
        limit_bytes=104857600,  # 100 MB default (STORE-01)
        used_bytes=0,
    )
    try:
        session.add(new_user)
        await session.flush()  # persist User before Quota FK
        session.add(quota)
        await session.commit()
        await session.refresh(new_user)
    except IntegrityError:
        await session.rollback()
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail="Email or handle already in use",
        )

    return {
        "id": str(new_user.id),
        "handle": new_user.handle,
        "email": new_user.email,
        "role": new_user.role,
        "totp_enabled": new_user.totp_enabled,
        "created_at": new_user.created_at.isoformat() if new_user.created_at else None,
    }


# ── POST /login ───────────────────────────────────────────────────────────────

@router.post("/login")
@limiter.limit("10/minute")
async def login(
    request: Request,
    body: LoginRequest,
    response: Response,
    session: AsyncSession = Depends(get_db),
):
    # Per-account rate limiting (SEC-02): Redis counter keyed by email, checked before any DB lookup
    redis_client = request.app.state.redis
    rate_key = f"login_attempts:{body.email}"
    count = await redis_client.incr(rate_key)
    if count == 1:
        await redis_client.expire(rate_key, 900)
    if count > 10:
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail="Too many login attempts. Try again in 15 minutes.",
        )

    result = await session.execute(select(User).where(User.email == str(body.email)))
    user: Optional[User] = result.scalar_one_or_none()

    _ip = get_client_ip(request)

    # Verify password (anti-enumeration: same error regardless of whether user exists)
    if user is None or not auth_service.verify_password(body.password, user.password_hash):
        await write_audit_log(
            session,
            event_type="auth.login_failed",
            user_id=user.id if user else None,
            actor_id=user.id if user else None,
            resource_id=None,
            ip_address=_ip,
            metadata_={"attempted_email_hash": hashlib.sha256(str(body.email).encode()).hexdigest()[:16]},
        )
        await session.commit()
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect email or password",
        )

    if not user.is_active:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Account deactivated",
        )

    # Password must change: return challenge without issuing tokens (T-02-16)
    if user.password_must_change:
        return {"requires_password_change": True, "user_id": str(user.id)}

    if user.totp_enabled:
        if body.totp_code is None and body.backup_code is None:
            return {"requires_totp": True}

        if body.totp_code is not None:
            # TOTP path takes precedence (even if backup_code also provided)
            ok = await auth_service.verify_totp(session, user.id, body.totp_code, redis_client)
            if not ok:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="Incorrect code",
                )
        else:
            # Backup code path (body.backup_code is not None and body.totp_code is None)
            ok = await auth_service.verify_backup_code(session, user.id, body.backup_code)
            if not ok:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="Invalid or already used code",
                )
            # D-13: backup code used event
            await write_audit_log(
                session,
                event_type="auth.backup_code_used",
                user_id=user.id,
                actor_id=user.id,
                resource_id=None,
                ip_address=_ip,
            )

    access_token = auth_service.create_access_token(
        str(user.id),
        user.role,
        user_agent=request.headers.get("User-Agent", ""),
        accept_lang=request.headers.get("Accept-Language", ""),
    )
    raw_refresh = await auth_service.create_refresh_token(session, user.id, remember_me=body.remember_me)
    _set_refresh_cookie(response, raw_refresh, remember_me=body.remember_me)

    # D-13: login success event
    await write_audit_log(
        session,
        event_type="auth.login",
        user_id=user.id,
        actor_id=user.id,
        resource_id=None,
        ip_address=_ip,
        metadata_={"totp_used": user.totp_enabled and body.totp_code is not None},
    )
    await session.commit()

    return {
        "access_token": access_token,
        "user": {
            "id": str(user.id),
            "handle": user.handle,
            "email": user.email,
            "role": user.role,
            "totp_enabled": user.totp_enabled,
        },
    }


# ── POST /refresh ─────────────────────────────────────────────────────────────

@router.post("/refresh")
@limiter.limit("10/minute")
async def refresh_token(
    request: Request,
    response: Response,
    session: AsyncSession = Depends(get_db),
):
    # Token reuse (revoked token presented) revokes entire family — family revocation security contract
    raw_token = request.cookies.get("refresh_token")
    if not raw_token:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="No refresh token",
        )

    try:
        new_raw, user_id_str = await auth_service.rotate_refresh_token(session, raw_token)
    except ValueError as exc:
        if "token_family_revoked" in str(exc):
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Session revoked",
            ) from exc
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or expired refresh token",
        ) from exc

    user = await session.get(User, uuid.UUID(user_id_str))
    if user is None or not user.is_active:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="User not found or deactivated",
        )

    _set_refresh_cookie(response, new_raw)

    access_token = auth_service.create_access_token(
        user_id_str,
        user.role,
        user_agent=request.headers.get("User-Agent", ""),
        accept_lang=request.headers.get("Accept-Language", ""),
    )
    return {
        "access_token": access_token,
        "user": {
            "id": str(user.id),
            "handle": user.handle,
            "email": user.email,
            "role": user.role,
            "totp_enabled": user.totp_enabled,
        },
    }


# ── POST /logout ──────────────────────────────────────────────────────────────

@router.post("/logout")
async def logout(request: Request, response: Response, session: AsyncSession = Depends(get_db)):
    import hashlib as _hashlib

    _ip = get_client_ip(request)

    raw_token = request.cookies.get("refresh_token")
    _logout_user_id = None
    if raw_token:
        token_hash = _hashlib.sha256(raw_token.encode()).hexdigest()
        result = await session.execute(
            select(RefreshToken).where(RefreshToken.token_hash == token_hash)
        )
        row: Optional[RefreshToken] = result.scalar_one_or_none()
        if row is not None:
            _logout_user_id = row.user_id
            row.revoked = True
            # D-13: logout event (written before commit, within same transaction)
            await write_audit_log(
                session,
                event_type="auth.logout",
                user_id=_logout_user_id,
                actor_id=_logout_user_id,
                resource_id=None,
                ip_address=_ip,
            )
            await session.commit()

    response.delete_cookie("refresh_token", path="/api/auth/refresh")
    return {"message": "Logged out"}


# ── POST /logout-all ──────────────────────────────────────────────────────────

@router.post("/logout-all")
async def logout_all(
    request: Request,
    response: Response,
    session: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _ip = get_client_ip(request)
    count = await auth_service.revoke_all_refresh_tokens(session, current_user.id)
    # D-13: sign-out-all event
    await write_audit_log(
        session,
        event_type="auth.sign_out_all",
        user_id=current_user.id,
        actor_id=current_user.id,
        resource_id=None,
        ip_address=_ip,
        metadata_={"sessions_revoked": count},
    )
    await session.commit()
    response.delete_cookie("refresh_token", path="/api/auth/refresh")
    return {"message": f"Signed out of {count} session(s)"}


# ── GET /me ───────────────────────────────────────────────────────────────────

@router.get("/me")
async def get_me(current_user: User = Depends(get_current_user)):
    return _user_dict(current_user)


# ── GET /me/quota ─────────────────────────────────────────────────────────────

@router.get("/me/quota")
async def get_my_quota(
    current_user: User = Depends(get_current_user),
    session: AsyncSession = Depends(get_db),
):
    q = await session.get(Quota, current_user.id)
    if q is None:
        raise HTTPException(status_code=404, detail="Quota not found")
    return {"used_bytes": q.used_bytes, "limit_bytes": q.limit_bytes}


# ── GET /me/preferences ───────────────────────────────────────────────────────

@router.get("/me/preferences")
async def get_my_preferences(
    current_user: User = Depends(get_current_user),
):
    # Falls back to 'in_app' if the column is absent (migration not yet run)
    try:
        pdf_open_mode = current_user.pdf_open_mode
    except AttributeError:
        pdf_open_mode = "in_app"
    return {"pdf_open_mode": pdf_open_mode}


# ── PATCH /me/preferences ─────────────────────────────────────────────────────

@router.patch("/me/preferences")
async def update_my_preferences(
    body: PreferencesUpdate,
    session: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    # Pydantic Literal["in_app", "new_tab"] enforces strict allowlist (T-04-05-05)
    user = await session.get(User, current_user.id)
    if user is None:
        raise HTTPException(status_code=404, detail="User not found")
    user.pdf_open_mode = body.pdf_open_mode
    session.add(user)
    await session.commit()
    return {"pdf_open_mode": user.pdf_open_mode}