maester_enterprise/security.py

from __future__ import annotations

import os
import secrets
import time
import uuid
from collections import defaultdict, deque
from collections.abc import Awaitable, Callable

from fastapi import Request, status
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse, Response

from .config import Settings


PUBLIC_PATHS = {
    "/",
    "/health",
    "/ready",
    "/version",
    "/openapi.json",
    "/studio",
}

PUBLIC_PREFIXES = (
    "/docs",
    "/redoc",
)


def _is_public_path(path: str, settings: Settings) -> bool:
    if path in PUBLIC_PATHS:
        return True
    if path.startswith("/studio/"):
        return True
    if settings.allow_public_docs and path.startswith(PUBLIC_PREFIXES):
        return True
    return False


def _has_valid_review_token(path: str) -> bool:
    prefix = "/automation/reviews/public/"
    if not path.startswith(prefix):
        return False
    token = path[len(prefix):].split("/", 1)[0]
    if not token:
        return False
    from .automation.operations import verify_review_token

    return verify_review_token(token) is not None


class RequestContextMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
        request_id = request.headers.get("X-Request-ID") or uuid.uuid4().hex
        request.state.request_id = request_id
        start = time.perf_counter()
        response = await call_next(request)
        response.headers["X-Request-ID"] = request_id
        response.headers["X-Response-Time-ms"] = str(round((time.perf_counter() - start) * 1000, 2))
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["X-Frame-Options"] = "DENY"
        response.headers["Referrer-Policy"] = "no-referrer"
        response.headers["Permissions-Policy"] = "camera=(), microphone=(), geolocation=()"
        return response


class BodyLimitMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, settings: Settings) -> None:
        super().__init__(app)
        self.settings = settings

    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
        content_length = request.headers.get("content-length")
        if content_length:
            try:
                length = int(content_length)
            except ValueError:
                return JSONResponse({"detail": "Invalid Content-Length"}, status_code=400)
            if length > self.settings.request_body_limit_bytes:
                return JSONResponse({"detail": "Request body too large"}, status_code=413)
        return await call_next(request)


class ApiKeyMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, settings: Settings) -> None:
        super().__init__(app)
        self.settings = settings

    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
        if _is_public_path(request.url.path, self.settings) or _has_valid_review_token(request.url.path):
            return await call_next(request)
        if not self.settings.api_key_required:
            return await call_next(request)
        provided = request.headers.get("X-API-Key") or request.query_params.get("api_key")
        if not self.settings.api_key or not provided or not secrets.compare_digest(provided, self.settings.api_key):
            return JSONResponse(
                {"detail": "Invalid or missing API key", "request_id": getattr(request.state, "request_id", None)},
                status_code=status.HTTP_401_UNAUTHORIZED,
            )
        return await call_next(request)


class RateLimitMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, settings: Settings) -> None:
        super().__init__(app)
        self.settings = settings
        self.hits: dict[str, deque[float]] = defaultdict(deque)

    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
        if _is_public_path(request.url.path, self.settings):
            return await call_next(request)
        now = time.time()
        window_start = now - 60
        key = request.headers.get("X-API-Key") or request.client.host if request.client else "unknown"
        bucket = self.hits[key]
        while bucket and bucket[0] < window_start:
            bucket.popleft()
        if len(bucket) >= self.settings.rate_limit_per_minute:
            return JSONResponse({"detail": "Rate limit exceeded"}, status_code=429)
        bucket.append(now)
        return await call_next(request)


def require_configured_security(settings: Settings) -> None:
    if settings.api_key_required and not settings.api_key:
        raise RuntimeError(
            "MAESTER_API_KEY is required in production. Set MAESTER_ENV=development and MAESTER_ALLOW_DEV_NO_API_KEY=true only for local development."
        )
    render_enabled = os.getenv("MAESTER_ENABLE_RENDER_ENGINE", "true").strip().lower() in {"1", "true", "yes", "on"}
    if render_enabled and settings.environment == "production":
        missing = [
            name
            for name in ("AVA2LON_SIGNING_SECRET", "BASYX_SIGNING_SECRET")
            if not os.getenv(name)
        ]
        if missing:
            raise RuntimeError(f"Missing production signing secret(s): {', '.join(missing)}")