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 re
import uuid
from typing import 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 services import auth as auth_service
from slowapi import Limiter
from slowapi.util import get_remote_address

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)

# ── Password strength validation ─────────────────────────────────────────────
_PASSWORD_DETAIL = (
    "Password must be at least 12 characters and include uppercase, "
    "lowercase, a number, and a special character."
)


def _validate_password_strength(password: str) -> bool:
    """Return True if password passes all strength rules (AUTH-01).

    Rules: min 12 chars, has uppercase, has lowercase, has digit, has special char.
    """
    if len(password) < 12:
        return False
    if not re.search(r"[A-Z]", password):
        return False
    if not re.search(r"[a-z]", password):
        return False
    if not re.search(r"[0-9]", password):
        return False
    if not re.search(r"[^A-Za-z0-9]", password):
        return False
    return True


# ── 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
    if not _validate_password_strength(body.password):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=_PASSWORD_DETAIL,
        )

    # 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)
    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()

    # 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):
        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",
                )

    # 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)

    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

    raw_token = request.cookies.get("refresh_token")
    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:
            row.revoked = True
            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."""
    count = await auth_service.revoke_all_refresh_tokens(session, current_user.id)
    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)


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

@router.post("/change-password")
async def change_password(
    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
    if not _validate_password_strength(body.new_password):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=_PASSWORD_DETAIL,
        )

    # Update password
    user = await session.get(User, current_user.id)
    user.password_hash = auth_service.hash_password(body.new_password)
    await session.commit()

    return {"message": "Password updated"}