kite/backend/main.py

import asyncio
import logging
import time
import uuid
from contextlib import asynccontextmanager

import structlog
from redis import asyncio as aioredis
from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from minio import Minio
from slowapi import _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware
from sqlalchemy import select, text
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response as StarletteResponse
from starlette.types import ASGIApp, Receive, Scope, Send

from api.auth import limiter as auth_limiter
from api.documents import router as documents_router
from api.topics import router as topics_router
from config import settings
from db.session import AsyncSessionLocal, engine
from services.ai_config import seed_system_settings_from_env
from services.logging import setup_logging
from services.rate_limiting import account_limiter


# ── CSP / Security headers middleware ────────────────────────────────────────

class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    """Add Content-Security-Policy, X-Frame-Options, and X-Content-Type-Options
    to every response (SEC-05, T-02-14).
    """

    async def dispatch(self, request: Request, call_next):
        response = await call_next(request)
        response.headers["Content-Security-Policy"] = (
            "default-src 'self'; "
            "script-src 'self'; "
            "style-src 'self' 'unsafe-inline'; "
            "img-src 'self' data:; "
            "frame-ancestors 'none'"
        )
        response.headers["X-Frame-Options"] = "DENY"
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        return response


# ── Origin validation middleware (SEC-01, T-02-11) ────────────────────────────

class OriginValidationMiddleware(BaseHTTPMiddleware):
    """Reject state-changing requests from Origins not in settings.cors_origins.

    For any non-idempotent method (not GET/HEAD/OPTIONS): if the Origin header
    is present and not in the allowed list, return 403.

    Placed BEFORE CORSMiddleware so it runs first (Starlette applies middleware
    in reverse insertion order — last added runs first).
    """

    async def dispatch(self, request: Request, call_next):
        if request.method not in {"GET", "HEAD", "OPTIONS"}:
            origin = request.headers.get("Origin")
            if origin is not None and origin not in settings.cors_origins:
                return StarletteResponse(content="Forbidden", status_code=403)
        return await call_next(request)


# ── Correlation ID Middleware ─────────────────────────────────────────────────

class CorrelationIDMiddleware:
    """Generate per-request UUID correlation ID; bind to structlog contextvars.

    Uses raw ASGI (NOT BaseHTTPMiddleware) to avoid response-body buffering
    issues with streaming responses (RESEARCH.md Anti-Patterns).

    Execution order (Starlette reverse-insertion):
      Registered LAST in app.add_middleware() so it runs FIRST on every request.

    Per-request:
      1. clear_contextvars() — Pitfall 2 guard; prevents context bleed between
         requests handled by the same worker.
      2. bind_contextvars(correlation_id, path, method) — available in all
         downstream log calls.
      3. Appends X-Correlation-ID to the response headers.
      4. After response: bind duration_ms for final log emission.
    """

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        correlation_id = str(uuid.uuid4())
        start_ns = time.perf_counter_ns()

        # Pitfall 2: MUST be first — clears any context from a prior request
        # on the same worker before binding new values.
        structlog.contextvars.clear_contextvars()
        structlog.contextvars.bind_contextvars(
            correlation_id=correlation_id,
            path=scope.get("path", ""),
            method=scope.get("method", ""),
        )

        _response_status: int = 0

        async def send_with_header(message: dict) -> None:
            nonlocal _response_status
            if message["type"] == "http.response.start":
                _response_status = message.get("status", 0)
                headers = list(message.get("headers", []))
                headers.append((b"x-correlation-id", correlation_id.encode()))
                message = {**message, "headers": headers}
            await send(message)

        await self.app(scope, receive, send_with_header)

        duration_ms = (time.perf_counter_ns() - start_ns) / 1_000_000
        structlog.contextvars.bind_contextvars(duration_ms=round(duration_ms, 2))
        structlog.get_logger("docuvault.access").info(
            "request_complete",
            status_code=_response_status,
        )


# ── ES256 startup rotation ────────────────────────────────────────────────────

async def _rotate_tokens_on_algorithm_change(session) -> None:
    """Idempotent ES256 migration (Phase 7.3 D-04/D-05).

    On every boot, compares the jwt_algorithm marker row in system_settings against
    'ES256'. If absent or different, bulk-revokes all active refresh tokens via raw
    SQL UPDATE and upserts the marker (with is_active=False so the AI provider loader
    never returns this row). No-op on second boot.
    """
    from db.models import SystemSettings  # local import to avoid circular deps

    stmt = select(SystemSettings).where(SystemSettings.provider_id == "jwt_algorithm")
    result = await session.execute(stmt)
    row = result.scalar_one_or_none()

    if row is not None and row.model_name == "ES256":
        return  # Already migrated — idempotent no-op

    # Bulk-revoke all active refresh tokens (single raw SQL — no Python iteration)
    await session.execute(
        text("UPDATE refresh_tokens SET revoked = true WHERE revoked = false")
    )
    logging.getLogger(__name__).info(
        "ES256 startup rotation: bulk-revoked all active refresh tokens"
    )

    if row is None:
        session.add(SystemSettings(
            id=uuid.uuid4(),
            provider_id="jwt_algorithm",
            model_name="ES256",
            context_chars=0,
            is_active=False,
            api_key_enc=None,
            base_url=None,
        ))
    else:
        row.model_name = "ES256"

    await session.commit()


# ── Lifespan ──────────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    """FastAPI lifespan: initialize MinIO, Redis, and admin bootstrap at startup.

    D-07: bucket auto-create ensures the docuvault bucket exists on every reboot.
    MinIO client stored on app.state.minio for use in the /health endpoint.
    Redis stored on app.state.redis for per-account rate limiting (SEC-02) and
    TOTP replay prevention (AUTH-08).
    Admin bootstrap (D-04): idempotent, runs only if no users exist.
    """
    # Initialize structured logging first — all subsequent log calls use the
    # configured renderer (JSON in production, console in dev).
    setup_logging(json_logs=settings.log_json, log_level=settings.log_level)

    # MinIO bucket initialization (RESEARCH.md Pattern 4)
    minio_client = Minio(
        settings.minio_endpoint,
        access_key=settings.minio_access_key,
        secret_key=settings.minio_secret_key,
        secure=False,
    )
    exists = await asyncio.to_thread(minio_client.bucket_exists, settings.minio_bucket)
    if not exists:
        await asyncio.to_thread(minio_client.make_bucket, settings.minio_bucket)
    app.state.minio = minio_client

    # Redis init for per-account rate limiting + TOTP replay prevention
    app.state.redis = await aioredis.from_url(settings.redis_url)

    # Admin bootstrap (D-04)
    from services.auth import bootstrap_admin  # noqa: PLC0415
    async with AsyncSessionLocal() as session:
        await bootstrap_admin(session)

    # AI provider seed (D-04 / Phase 7): populate system_settings from env vars
    # on first boot.  Wrapped in try/except so that a missing table (fresh container
    # before migrations run) does not crash startup — logs a warning and skips.
    try:
        async with AsyncSessionLocal() as session:
            await seed_system_settings_from_env(session)
            await session.commit()
    except Exception as _seed_exc:
        import logging as _logging
        _logging.getLogger(__name__).warning(
            "AI provider seed skipped (table may not exist yet): %s", _seed_exc
        )

    # ES256 startup rotation (Phase 7.3 — D-04):
    # If jwt_algorithm in system_settings differs from "ES256" (or is absent),
    # bulk-revoke all refresh tokens and record the new algorithm.
    # Wrapped in try/except so a missing table before migrations doesn't crash.
    try:
        async with AsyncSessionLocal() as session:
            await _rotate_tokens_on_algorithm_change(session)
    except Exception as _es256_exc:
        logging.getLogger(__name__).warning(
            "ES256 rotation check skipped (table may not exist yet): %s", _es256_exc
        )

    yield

    # Shutdown: close pooled connections and Redis
    await app.state.redis.close()
    await engine.dispose()


# ── Application factory ───────────────────────────────────────────────────────

app = FastAPI(title="Document Scanner API", version="0.2.1", lifespan=lifespan)

# Rate limiter state (slowapi)
app.state.limiter = auth_limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
app.add_middleware(SlowAPIMiddleware)

# ── Middleware registration order (Starlette: last added = first to run) ───────
# Desired execution order (request path): Origin → CORS → SecurityHeaders → route
# Insertion order (last registered = first to run): SecurityHeaders → CORS → Origin
# Result: register SecurityHeaders first, then CORS, then Origin last.

# 1. Security headers (CSP etc.) — runs last in the chain
app.add_middleware(SecurityHeadersMiddleware)

# 2. CORS — updated to use settings.cors_origins (D-09); wildcard removed (T-02-15)
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.cors_origins,
    allow_credentials=True,   # Required for httpOnly cookie flow
    allow_methods=["*"],
    allow_headers=["*"],
)

# 3. Origin validation — runs second (added second-to-last), before CORS and route handlers
app.add_middleware(OriginValidationMiddleware)

# 4. CorrelationID — added LAST so it runs FIRST (Starlette reverse-insertion order)
app.add_middleware(CorrelationIDMiddleware)


# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/health")
async def health(request: Request):
    """Extended health probe: reports PostgreSQL and MinIO connectivity (D-07).

    Always returns HTTP 200 — 'degraded' status signals a partial outage without
    causing load-balancer retries.

    Note (T-01-05-03): error strings expose Python exception class names — acceptable
    for an internal/dev endpoint in Phase 1. Phase 2 will trim to 'error' or
    'unhealthy' once the endpoint is internet-facing.
    """
    checks: dict = {}

    # PostgreSQL probe
    try:
        async with AsyncSessionLocal() as session:
            await session.execute(text("SELECT 1"))
        checks["postgres"] = "ok"
    except Exception as e:
        checks["postgres"] = f"error: {type(e).__name__}: {e}"

    # MinIO probe
    try:
        ok = await asyncio.to_thread(
            request.app.state.minio.bucket_exists, settings.minio_bucket
        )
        checks["minio"] = "ok" if ok else "error: bucket missing"
    except Exception as e:
        checks["minio"] = f"error: {type(e).__name__}: {e}"

    status_val = "ok" if all(v == "ok" for v in checks.values()) else "degraded"
    return {"status": status_val, "checks": checks}


# ── Include routers ───────────────────────────────────────────────────────────

app.include_router(documents_router)
app.include_router(topics_router)

# Phase 2: auth and admin routers
from api.auth import router as auth_router  # noqa: E402
from api.admin import router as admin_router  # noqa: E402
app.include_router(auth_router)
app.include_router(admin_router)

# Phase 4: folders router (FOLD-01..05) and document-move endpoint
from api.folders import router as folders_router  # noqa: E402
from api.folders import document_move_router as document_move_router  # noqa: E402
app.include_router(folders_router)
app.include_router(document_move_router)

# Phase 4: shares router (SHARE-01..05)
from api.shares import router as shares_router  # noqa: E402
app.include_router(shares_router)

# Phase 4: audit log viewer + CSV export (ADMIN-06)
from api.audit import router as audit_router  # noqa: E402
app.include_router(audit_router)

# Phase 5: cloud storage backend connection management (CLOUD-01..07)
from api.cloud import router as cloud_router, users_router as cloud_users_router  # noqa: E402
app.include_router(cloud_router)
app.include_router(cloud_users_router)