kite/backend/api/auth/totp.py

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),
):
    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),
):
    # Backup codes are returned here ONLY — stored as Argon2 hashes, 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",
        )

    user = await session.get(User, current_user.id)
    user.totp_enabled = True
    await session.flush()

    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),
):
    _ip = get_client_ip(request)
    user = await session.get(User, current_user.id)
    user.totp_enabled = False
    user.totp_secret = None

    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}