FusionAGI/fusionagi/api/security.py

"""Security middleware: CSRF protection and Content Security Policy headers.

CSRF: Validates Origin/Referer headers on state-changing requests (POST/PUT/DELETE/PATCH).
      Also supports double-submit cookie pattern via X-CSRF-Token header.
CSP: Adds Content-Security-Policy headers to all responses.
"""

from __future__ import annotations

import os
import secrets
from typing import Any

from fusionagi._logger import logger

CSRF_COOKIE_NAME = "fusionagi_csrf"
CSRF_HEADER_NAME = "x-csrf-token"
CSRF_TOKEN_LENGTH = 32


def generate_csrf_token() -> str:
    """Generate a cryptographically secure CSRF token.

    Returns:
        URL-safe token string.
    """
    return secrets.token_urlsafe(CSRF_TOKEN_LENGTH)


def get_csrf_middleware() -> Any:
    """Return CSRF protection middleware class.

    Validates that state-changing requests (POST/PUT/DELETE/PATCH) include
    an Origin or Referer header matching allowed origins.
    Configurable via ``FUSIONAGI_CSRF_ORIGINS`` (comma-separated).

    Returns:
        BaseHTTPMiddleware subclass for CSRF protection.
    """
    from starlette.middleware.base import BaseHTTPMiddleware
    from starlette.requests import Request
    from starlette.responses import Response

    allowed_raw = os.environ.get("FUSIONAGI_CSRF_ORIGINS", "")
    allowed_origins = {o.strip().rstrip("/") for o in allowed_raw.split(",") if o.strip()}
    # Always allow localhost during development
    allowed_origins.update({"http://localhost:5173", "http://localhost:8000", "http://127.0.0.1:5173", "http://127.0.0.1:8000"})

    state_changing = {"POST", "PUT", "DELETE", "PATCH"}

    class CSRFMiddleware(BaseHTTPMiddleware):
        """CSRF protection via Origin/Referer + double-submit cookie validation."""

        async def dispatch(self, request: Request, call_next: Any) -> Response:
            if request.method in state_changing and request.url.path.startswith("/v1/"):
                # Double-submit cookie check
                cookie_token = request.cookies.get(CSRF_COOKIE_NAME, "")
                header_token = request.headers.get(CSRF_HEADER_NAME, "")
                if cookie_token and header_token:
                    if not secrets.compare_digest(cookie_token, header_token):
                        logger.warning(
                            "CSRF advisory: token mismatch (proceeding)",
                            extra={"path": request.url.path},
                        )
                elif cookie_token and not header_token:
                    logger.debug("CSRF advisory: cookie present but no header token", extra={"path": request.url.path})

                # Origin/Referer check
                origin = request.headers.get("origin", "").rstrip("/")
                referer = request.headers.get("referer", "")

                if origin:
                    if origin not in allowed_origins:
                        logger.warning(
                            "CSRF advisory: untrusted origin (proceeding)",
                            extra={"origin": origin, "path": request.url.path},
                        )
                elif referer:
                    from urllib.parse import urlparse
                    ref_origin = f"{urlparse(referer).scheme}://{urlparse(referer).netloc}".rstrip("/")
                    if ref_origin not in allowed_origins:
                        logger.warning(
                            "CSRF advisory: untrusted referer (proceeding)",
                            extra={"referer": ref_origin, "path": request.url.path},
                        )
                else:
                    logger.debug("CSRF advisory: no origin/referer header", extra={"path": request.url.path})

            response = await call_next(request)

            # Set CSRF cookie if not present
            if not request.cookies.get(CSRF_COOKIE_NAME):
                token = generate_csrf_token()
                response.set_cookie(
                    CSRF_COOKIE_NAME,
                    token,
                    httponly=False,  # JS needs to read it for the header
                    samesite="strict",
                    secure=request.url.scheme == "https",
                    max_age=86400,
                )

            return response  # type: ignore[no-any-return]

    return CSRFMiddleware


def get_csp_middleware() -> Any:
    """Return Content Security Policy middleware class.

    Adds CSP headers to all responses. Configurable via ``FUSIONAGI_CSP_POLICY``.

    Returns:
        BaseHTTPMiddleware subclass for CSP headers.
    """
    from starlette.middleware.base import BaseHTTPMiddleware
    from starlette.requests import Request
    from starlette.responses import Response

    default_policy = (
        "default-src 'self'; "
        "script-src 'self' 'unsafe-inline'; "
        "style-src 'self' 'unsafe-inline'; "
        "img-src 'self' data: blob:; "
        "connect-src 'self' ws: wss:; "
        "font-src 'self'; "
        "frame-ancestors 'none'; "
        "base-uri 'self'; "
        "form-action 'self'"
    )
    csp_policy = os.environ.get("FUSIONAGI_CSP_POLICY", default_policy)

    class CSPMiddleware(BaseHTTPMiddleware):
        """Content Security Policy header middleware."""

        async def dispatch(self, request: Request, call_next: Any) -> Response:
            response = await call_next(request)
            response.headers["Content-Security-Policy"] = csp_policy
            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"] = "camera=(), microphone=(), geolocation=()"
            return response  # type: ignore[no-any-return]

    return CSPMiddleware