kite/backend/api/auth/totp.py

"""
Auth API — TOTP endpoints.

Handles:
  GET    /totp/setup   — provision TOTP secret for current user
  POST   /totp/enable  — enable TOTP (verify code, generate backup codes)
  DELETE /totp         — disable TOTP (clear secret, revoke other sessions)
"""
from __future__ import annotations

import hashlib
import time

from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from sqlalchemy import delete
from sqlalchemy.ext.asyncio import AsyncSession

from config import settings
from db.models import BackupCode, 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, TotpEnableRequest

router = APIRouter()


# ── GET /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 /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)

    # Revoke other sessions; keep current one alive via skip_token_hash (CR-02)
    raw_cookie = request.cookies.get("refresh_token")
    skip_hash = hashlib.sha256(raw_cookie.encode()).hexdigest() if raw_cookie else None
    revoked = await auth_service.revoke_all_refresh_tokens(session, current_user.id, skip_token_hash=skip_hash)

    # 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,
        metadata_={"sessions_revoked": revoked},
    )
    # Revoke any pre-enroll access tokens still within their TTL window (T-7.2-01)
    await redis_client.set(
        f"user_nbf:{current_user.id}",
        int(time.time()),
        ex=settings.access_token_expire_minutes * 60,
    )
    await session.commit()

    return {"backup_codes": plain_codes, "sessions_revoked": revoked}


# ── DELETE /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))

    # Revoke other sessions; keep current one alive via skip_token_hash (CR-03)
    raw_cookie = request.cookies.get("refresh_token")
    skip_hash = hashlib.sha256(raw_cookie.encode()).hexdigest() if raw_cookie else None
    revoked = await auth_service.revoke_all_refresh_tokens(session, current_user.id, skip_token_hash=skip_hash)

    # 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,
        metadata_={"sessions_revoked": revoked},
    )
    # Revoke any pre-revoke access tokens still within their TTL window (T-7.2-01)
    await request.app.state.redis.set(
        f"user_nbf:{current_user.id}",
        int(time.time()),
        ex=settings.access_token_expire_minutes * 60,
    )
    await session.commit()

    return {"message": "TOTP disabled", "sessions_revoked": revoked}