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.6", 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)