Files
agent_delta/backend/app/core/middleware.py
T
2026-06-15 14:22:07 +08:00

103 lines
3.6 KiB
Python

"""Application middleware."""
from typing import ClassVar
from uuid import uuid4
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp
class RequestIDMiddleware(BaseHTTPMiddleware):
"""Middleware that adds a unique request ID to each request.
The request ID is taken from the X-Request-ID header if present,
otherwise a new UUID is generated. The ID is added to the response
headers and is available in request.state.request_id.
"""
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
"""Add request ID to request state and response headers."""
request_id = request.headers.get("X-Request-ID", str(uuid4()))
request.state.request_id = request_id
response = await call_next(request)
response.headers["X-Request-ID"] = request_id
return response
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
"""Middleware that adds security headers to all responses.
This includes:
- Content-Security-Policy (CSP)
- X-Content-Type-Options
- X-Frame-Options
- X-XSS-Protection
- Referrer-Policy
- Permissions-Policy
Usage:
app.add_middleware(SecurityHeadersMiddleware)
# Or with custom CSP:
app.add_middleware(
SecurityHeadersMiddleware,
csp_directives={
"default-src": "'self'",
"script-src": "'self' 'unsafe-inline'",
}
)
"""
DEFAULT_CSP_DIRECTIVES: ClassVar[dict[str, str]] = {
"default-src": "'self'",
"script-src": "'self'",
"style-src": "'self' 'unsafe-inline'", # Allow inline styles for some UI libs
"img-src": "'self' data: https:",
"font-src": "'self' data:",
"connect-src": "'self'",
"frame-ancestors": "'none'",
"base-uri": "'self'",
"form-action": "'self'",
}
def __init__(
self,
app: ASGIApp,
csp_directives: dict[str, str] | None = None,
exclude_paths: set[str] | None = None,
) -> None:
super().__init__(app)
self.csp_directives = csp_directives or self.DEFAULT_CSP_DIRECTIVES
self.exclude_paths = exclude_paths or {"/docs", "/redoc", "/openapi.json"}
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
"""Add security headers to the response."""
response = await call_next(request)
# Skip for docs/openapi endpoints which need different CSP
if request.url.path in self.exclude_paths:
return response
# Build CSP header
csp_value = "; ".join(
f"{directive} {value}" for directive, value in self.csp_directives.items()
)
# Add security headers — respect any already set by the route, so an
# endpoint can opt into less restrictive framing (e.g. user-content
# files served inline for the chat preview panel).
response.headers.setdefault("Content-Security-Policy", csp_value)
response.headers.setdefault("X-Content-Type-Options", "nosniff")
response.headers.setdefault("X-Frame-Options", "DENY")
response.headers.setdefault("X-XSS-Protection", "1; mode=block")
response.headers.setdefault("Referrer-Policy", "strict-origin-when-cross-origin")
response.headers["Permissions-Policy"] = (
"accelerometer=(), camera=(), geolocation=(), gyroscope=(), "
"magnetometer=(), microphone=(), payment=(), usb=()"
)
return response