"""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