app/security/middleware.py

"""Security middleware: headers, body size, injection scan on JSON bodies."""
import json
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

from app.security.input_guard import MAX_TEXT_FIELD, _INJECTION, _SSTI

MAX_BODY_BYTES = 512_000  # 512 KB


class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next) -> Response:
        response = await call_next(request)
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["X-Frame-Options"] = "DENY"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        response.headers["Permissions-Policy"] = "geolocation=(), microphone=(), camera=()"
        if request.url.path.startswith("/api/auth"):
            response.headers["Cache-Control"] = "no-store"
        return response


class RequestGuardMiddleware(BaseHTTPMiddleware):
    """Block oversized bodies and scan auth JSON for SSTI/injection probes."""

    async def dispatch(self, request: Request, call_next) -> Response:
        if request.method in ("POST", "PUT", "PATCH") and request.url.path.startswith("/api/auth"):
            content_length = request.headers.get("content-length")
            if content_length and int(content_length) > MAX_BODY_BYTES:
                return JSONResponse(status_code=413, content={"detail": "Request body too large"})

            ctype = (request.headers.get("content-type") or "").lower()
            if "application/json" in ctype:
                body = await request.body()
                if len(body) > MAX_BODY_BYTES:
                    return JSONResponse(status_code=413, content={"detail": "Request body too large"})
                if body:
                    try:
                        data = json.loads(body)
                        if _json_has_threats(data):
                            return JSONResponse(status_code=400, content={"detail": "Invalid request content blocked"})
                    except json.JSONDecodeError:
                        return JSONResponse(status_code=400, content={"detail": "Invalid JSON"})

                async def receive():
                    return {"type": "http.request", "body": body, "more_body": False}

                request = Request(request.scope, receive)

        return await call_next(request)


def _json_has_threats(obj, depth: int = 0) -> bool:
    if depth > 12:
        return True
    if isinstance(obj, str):
        if len(obj) > MAX_TEXT_FIELD:
            return True
        if _SSTI.search(obj) or _INJECTION.search(obj):
            return True
        return False
    if isinstance(obj, dict):
        return any(_json_has_threats(v, depth + 1) for v in obj.values())
    if isinstance(obj, list):
        return any(_json_has_threats(v, depth + 1) for v in obj[:200])
    return False