mirror of
https://github.com/furyhawk/agent_delta.git
synced 2026-07-21 02:05:36 +00:00
103 lines
3.6 KiB
Python
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
|