Adriano Dal Pastro
98ad65d248
Adds middleware requiring X-Api-Key header or ?api_key= query param to access Swagger UI and ReDoc. OpenAPI schema remains public. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
import asyncio
|
|
import logging
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import JSONResponse
|
|
from slowapi import _rate_limit_exceeded_handler
|
|
from slowapi.errors import RateLimitExceeded
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.requests import Request
|
|
|
|
from src.backend.api.dependencies import limiter
|
|
from src.backend.config import settings
|
|
from src.backend.api.routes.health import router as health_router
|
|
from src.backend.api.routes.requests import router as requests_router
|
|
from src.backend.api.routes.sessions import router as sessions_router
|
|
from src.backend.api.routes.topics import router as topics_router
|
|
from src.backend.db.database import init_db
|
|
from src.backend.services.queue import queue_loop
|
|
from src.backend.services.worker_client import WorkerClient
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
await init_db()
|
|
worker = WorkerClient()
|
|
task = asyncio.create_task(queue_loop(worker))
|
|
yield
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
|
|
app = FastAPI(
|
|
title="OpusAgent",
|
|
version="0.1.0",
|
|
docs_url="/docs",
|
|
redoc_url="/redoc",
|
|
openapi_url="/openapi.json",
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
PROTECTED_DOCS_PATHS = {"/docs", "/redoc"}
|
|
|
|
|
|
class DocsAuthMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(self, request: Request, call_next):
|
|
if request.url.path in PROTECTED_DOCS_PATHS:
|
|
key = request.headers.get("X-Api-Key") or request.query_params.get("api_key")
|
|
if key != settings.api_key:
|
|
return JSONResponse(status_code=403, content={"detail": "API key required"})
|
|
return await call_next(request)
|
|
|
|
|
|
app.add_middleware(DocsAuthMiddleware)
|
|
app.state.limiter = limiter
|
|
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
app.include_router(health_router, prefix="/api")
|
|
app.include_router(topics_router, prefix="/api")
|
|
app.include_router(requests_router, prefix="/api")
|
|
app.include_router(sessions_router, prefix="/api")
|