kite/backend/api/auth.py

"""
Auth API endpoints for DocuVault.

Implements:
  POST /api/auth/register        — new user registration with HIBP check
  POST /api/auth/login           — login with optional TOTP/backup-code second factor
  POST /api/auth/refresh         — rotate refresh token (httpOnly cookie in/out)
  POST /api/auth/logout          — revoke current refresh token, clear cookie
  GET  /api/auth/me              — return current user profile
  POST /api/auth/change-password — update password (requires current password)

Security invariants:
  - Per-account rate limit: 10 login attempts per email per 15 minutes (SEC-02)
  - HTTP 429 returned before any DB lookup when the counter is exceeded
  - httpOnly Secure SameSite=Strict refresh cookie (CLAUDE.md constraint)
  - HIBP breach check on register and change-password (SEC-03)
  - TOTP takes precedence over backup_code when both fields are provided
  - password_must_change=True: returns requires_password_change without tokens
"""
from __future__ import annotations

import uuid
from typing import Literal, Optional

from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from pydantic import BaseModel, EmailStr
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from config import settings
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 slowapi import Limiter
from slowapi.util import get_remote_address
from sqlalchemy import delete

router = APIRouter(prefix="/api/auth", tags=["auth"])

# IP-level rate limiter (SEC-02 — 10 req/min on register/login/refresh)
limiter = Limiter(key_func=get_remote_address)


# ── Request models ────────────────────────────────────────────────────────────

class RegisterRequest(BaseModel):
    handle: str
    email: EmailStr
    password: str


class LoginRequest(BaseModel):
    email: EmailStr
    password: str
    totp_code: Optional[str] = None
    backup_code: Optional[str] = None


class ChangePasswordRequest(BaseModel):
    current_password: str
    new_password: str


# ── Helper: set httpOnly refresh cookie ──────────────────────────────────────

def _set_refresh_cookie(response: Response, raw_token: str) -> None:
    """Set the httpOnly Secure SameSite=Strict refresh cookie (CLAUDE.md constraint)."""
    response.set_cookie(
        key="refresh_token",
        value=raw_token,
        httponly=True,
        secure=True,
        samesite="strict",
        path="/api/auth/refresh",
        max_age=settings.refresh_token_expire_days * 86400,
    )


def _user_dict(user: User) -> dict:
    """Return serialisable user metadata (no password_hash, no credentials_enc)."""
    return {
        "id": str(user.id),
        "handle": user.handle,
        "email": user.email,
        "role": user.role,
        "totp_enabled": user.totp_enabled,
        "created_at": user.created_at.isoformat() if user.created_at else None,
    }


# ── POST /api/auth/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),
):
    """Register a new user account.

    - Validates password strength (min 12 chars, upper, lower, digit, special)
    - Checks HIBP k-anonymity API for breached passwords
    - Hashes password with Argon2
    - Inserts User + Quota rows in a single transaction
    """
    # Password strength check
    try:
        auth_service.validate_password_strength(body.password)
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc))

    # HIBP breach check
    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.",
        )

    # Duplicate email/handle check
    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",
        )

    # Create user and quota
    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,
    )
    session.add(new_user)
    await session.flush()  # persist User before Quota FK
    session.add(quota)
    await session.commit()
    await session.refresh(new_user)

    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 /api/auth/login ──────────────────────────────────────────────────────

@router.post("/login")
@limiter.limit("10/minute")
async def login(
    request: Request,
    body: LoginRequest,
    response: Response,
    session: AsyncSession = Depends(get_db),
):
    """Authenticate a user and issue tokens.

    Per-account rate limiting (SEC-02): checks Redis counter keyed by email
    BEFORE any DB lookup to prevent enumeration timing attacks.

    Three login flows:
      1. No TOTP enabled: password → tokens
      2. TOTP enabled, no code provided: requires_totp = True (challenge)
      3. TOTP enabled, totp_code provided: verify TOTP → tokens
      4. TOTP enabled, backup_code provided (no totp_code): verify backup → tokens
    """
    # Per-account rate limiting (SEC-02)
    redis_client = request.app.state.redis
    rate_key = f"login_attempts:{body.email}"
    count = await redis_client.incr(rate_key)
    if count == 1:
        # Set TTL only on first increment (15-minute window)
        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.",
        )

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

    # IP extraction for audit log (used in both success and failure paths)
    _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": str(body.email)},
        )
        await session.commit()
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect email or password",
        )

    # Active check
    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)}

    # TOTP second-factor dispatch
    if user.totp_enabled:
        if body.totp_code is None and body.backup_code is None:
            # Challenge: prompt for second factor
            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,
            )

    # Issue tokens
    access_token = auth_service.create_access_token(str(user.id), user.role)
    raw_refresh = await auth_service.create_refresh_token(session, user.id)
    _set_refresh_cookie(response, raw_refresh)

    # 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 /api/auth/refresh ────────────────────────────────────────────────────

@router.post("/refresh")
@limiter.limit("10/minute")
async def refresh_token(
    request: Request,
    response: Response,
    session: AsyncSession = Depends(get_db),
):
    """Rotate the refresh token.

    Reads the refresh_token httpOnly cookie; on success issues a new access
    token and rotates the refresh cookie.
    On token reuse (revoked token presented), revokes entire family and raises 401.
    """
    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

    # Look up user for response body
    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 new refresh cookie
    _set_refresh_cookie(response, new_raw)

    access_token = auth_service.create_access_token(user_id_str, user.role)
    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 /api/auth/logout ─────────────────────────────────────────────────────

@router.post("/logout")
async def logout(request: Request, response: Response, session: AsyncSession = Depends(get_db)):
    """Revoke current refresh token and clear the cookie."""
    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 /api/auth/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),
):
    """Sign out of all devices: revoke all refresh tokens for 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 /api/auth/me ──────────────────────────────────────────────────────────

@router.get("/me")
async def get_me(current_user: User = Depends(get_current_user)):
    """Return the current user's profile (requires valid Bearer token)."""
    return _user_dict(current_user)


# ── GET /api/auth/me/quota ────────────────────────────────────────────────────

@router.get("/me/quota")
async def get_my_quota(
    current_user: User = Depends(get_current_user),
    session: AsyncSession = Depends(get_db),
):
    """Return the current user's quota usage (STORE-04).

    Returns {"used_bytes": int, "limit_bytes": int} for the sidebar quota bar.
    Quota row is created at registration (100 MB default — STORE-01).
    """
    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}


# ── POST /api/auth/change-password ───────────────────────────────────────────

@router.post("/change-password")
async def change_password(
    request: Request,
    body: ChangePasswordRequest,
    session: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """Update the current user's password.

    Checks:
      1. current_password matches stored hash
      2. new_password has not appeared in HIBP (SEC-03)
      3. new_password meets strength requirements (AUTH-01)
    """
    # Verify current password
    if not auth_service.verify_password(body.current_password, current_user.password_hash):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Current password is incorrect",
        )

    # HIBP breach check on new password (SEC-03)
    if await auth_service.check_hibp(body.new_password):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="This password has appeared in a data breach. Choose a different password.",
        )

    # Password strength check
    try:
        auth_service.validate_password_strength(body.new_password)
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc))

    # Update password
    _ip = get_client_ip(request)
    user = await session.get(User, current_user.id)
    user.password_hash = auth_service.hash_password(body.new_password)
    # D-13: password changed event (flush within same transaction before commit)
    await write_audit_log(
        session,
        event_type="auth.password_changed",
        user_id=current_user.id,
        actor_id=current_user.id,
        resource_id=None,
        ip_address=_ip,
    )
    await session.commit()

    return {"message": "Password updated"}


# ── Request models for new endpoints ─────────────────────────────────────────

class TotpEnableRequest(BaseModel):
    code: str


class PasswordResetRequest(BaseModel):
    email: EmailStr


class PasswordResetConfirmRequest(BaseModel):
    token: str
    new_password: str


# ── GET /api/auth/totp/setup ──────────────────────────────────────────────────

@router.get("/totp/setup")
async def totp_setup(
    session: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """Provision a TOTP secret for the current user.

    If TOTP is already enabled, returns 400.
    Returns { provisioning_uri, secret } — the provisioning_uri is suitable
    for QR code generation. The secret is the base32-encoded TOTP secret.
    """
    if current_user.totp_enabled:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="TOTP already enabled",
        )
    secret, provisioning_uri = await auth_service.provision_totp(session, current_user.id)
    return {"provisioning_uri": provisioning_uri, "secret": secret}


# ── POST /api/auth/totp/enable ───────────────────────────────────────────────

@router.post("/totp/enable")
@limiter.limit("10/minute")
async def enable_totp(
    request: Request,
    body: TotpEnableRequest,
    session: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """Enable TOTP for the current user.

    Rate-limited to 10 attempts/minute per IP (SEC-02 / T-02-25).
    Verifies the submitted 6-digit code (with Redis replay prevention, AUTH-08).
    On success: marks TOTP enabled, generates and returns 10 one-time backup codes.
    The backup codes are ONLY returned here — they are stored as Argon2 hashes
    in the DB and never returned again (T-02-19).
    """
    redis_client = request.app.state.redis
    ok = await auth_service.verify_totp(session, current_user.id, body.code, redis_client)
    if not ok:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Incorrect or expired code",
        )

    # Mark TOTP as enabled
    user = await session.get(User, current_user.id)
    user.totp_enabled = True
    await session.flush()

    # Generate and store 10 backup codes; return plaintext to user (one-time, T-02-19)
    plain_codes = auth_service.generate_backup_codes(10)
    await auth_service.store_backup_codes(session, current_user.id, plain_codes)

    # D-13: TOTP enrolled event
    _ip = get_client_ip(request)
    await write_audit_log(
        session,
        event_type="auth.totp_enrolled",
        user_id=current_user.id,
        actor_id=current_user.id,
        resource_id=None,
        ip_address=_ip,
    )
    await session.commit()

    return {"backup_codes": plain_codes}


# ── DELETE /api/auth/totp ─────────────────────────────────────────────────────

@router.delete("/totp")
async def disable_totp(
    request: Request,
    session: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """Disable TOTP for the current user.

    Clears totp_secret, sets totp_enabled=False, and deletes all backup codes.
    """
    _ip = get_client_ip(request)
    user = await session.get(User, current_user.id)
    user.totp_enabled = False
    user.totp_secret = None

    # Delete all backup codes for this user (including unused ones)
    await session.execute(delete(BackupCode).where(BackupCode.user_id == current_user.id))

    # D-13: TOTP revoked event
    await write_audit_log(
        session,
        event_type="auth.totp_revoked",
        user_id=current_user.id,
        actor_id=current_user.id,
        resource_id=None,
        ip_address=_ip,
    )
    await session.commit()

    return {"message": "TOTP disabled"}


# ── POST /api/auth/password-reset ─────────────────────────────────────────────

@router.post("/password-reset", status_code=status.HTTP_202_ACCEPTED)
@limiter.limit("5/hour")
async def password_reset_request(
    request: Request,
    body: PasswordResetRequest,
    session: AsyncSession = Depends(get_db),
):
    """Request a password reset email.

    Always returns 202 regardless of whether the email exists (anti-enumeration, T-02-22).
    If the user is found, a signed reset token (1-hour JWT) is generated and a Celery
    task is enqueued to send the email (D-02, D-03).
    """
    from sqlalchemy import select as _select  # noqa: PLC0415 (already imported above)

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

    if user is not None:
        token = auth_service.create_password_reset_token(str(user.id))
        reset_link = f"{settings.frontend_url}/password-reset/confirm?token={token}"
        # Deferred import to avoid circular import; Celery task is fire-and-forget
        from tasks.email_tasks import send_reset_email  # noqa: PLC0415
        send_reset_email.delay(user.email, reset_link)

    # Always return 202 (anti-enumeration — never reveal whether email exists)
    return {
        "message": (
            "If an account exists for that email, you will receive a reset link shortly."
        )
    }


# ── POST /api/auth/password-reset/confirm ────────────────────────────────────

@router.post("/password-reset/confirm")
async def password_reset_confirm(
    body: PasswordResetConfirmRequest,
    session: AsyncSession = Depends(get_db),
):
    """Confirm a password reset using the token from the email link.

    Validates the reset token, enforces password strength + HIBP check, updates
    the password, and revokes all refresh tokens. Does NOT issue new tokens —
    the user must sign in again through /login (AUTH-05, T-02-21).
    """
    try:
        user_id_str = auth_service.decode_password_reset_token(body.token)
    except ValueError:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid or expired reset link",
        )

    # Password strength validation
    try:
        auth_service.validate_password_strength(body.new_password)
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc))

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

    # Load user
    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_400_BAD_REQUEST,
            detail="Invalid or expired reset link",
        )

    # Update password and revoke all sessions (forces re-auth through TOTP if enabled)
    user.password_hash = auth_service.hash_password(body.new_password)
    await auth_service.revoke_all_refresh_tokens(session, user.id)
    await session.commit()

    # Do NOT issue tokens (AUTH-05 — user must pass TOTP gate on next login)
    return {"message": "Password updated. Please sign in."}


# ── Preferences models ────────────────────────────────────────────────────────

class PreferencesUpdate(BaseModel):
    """Request body for PATCH /api/auth/me/preferences.

    Validates pdf_open_mode strictly via Literal (T-04-05-05 — no mass assignment).
    """
    pdf_open_mode: Literal["in_app", "new_tab"]


# ── GET /api/auth/me/preferences ─────────────────────────────────────────────

@router.get("/me/preferences")
async def get_my_preferences(
    current_user: User = Depends(get_current_user),
):
    """Return the current user's PDF open mode preference (D-10).

    Both regular users and admins can read their own preferences.
    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 /api/auth/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),
):
    """Update the current user's PDF open mode preference (D-10).

    Both regular users and admins can update their own preferences.
    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}