breakpilot-compliance/backend-compliance/middleware/security_headers.py

"""
Security Headers Middleware

Adds security headers to all HTTP responses to protect against common attacks.

Headers added:
- X-Content-Type-Options: nosniff
- X-Frame-Options: DENY
- X-XSS-Protection: 1; mode=block
- Strict-Transport-Security (HSTS)
- Content-Security-Policy
- Referrer-Policy
- Permissions-Policy

Usage:
    from middleware import SecurityHeadersMiddleware

    app.add_middleware(SecurityHeadersMiddleware)

    # Or with custom configuration:
    app.add_middleware(
        SecurityHeadersMiddleware,
        hsts_enabled=True,
        csp_policy="default-src 'self'",
    )
"""

import os
from dataclasses import dataclass, field
from typing import Dict, List, Optional

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response


@dataclass
class SecurityHeadersConfig:
    """Configuration for security headers."""

    # X-Content-Type-Options
    content_type_options: str = "nosniff"

    # X-Frame-Options
    frame_options: str = "DENY"

    # X-XSS-Protection (legacy, but still useful for older browsers)
    xss_protection: str = "1; mode=block"

    # Strict-Transport-Security
    hsts_enabled: bool = True
    hsts_max_age: int = 31536000  # 1 year
    hsts_include_subdomains: bool = True
    hsts_preload: bool = False

    # Content-Security-Policy
    csp_enabled: bool = True
    csp_policy: str = "default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'; img-src 'self' data: https:; font-src 'self' data:; connect-src 'self' https:; frame-ancestors 'none'"

    # Referrer-Policy
    referrer_policy: str = "strict-origin-when-cross-origin"

    # Permissions-Policy (formerly Feature-Policy)
    permissions_policy: str = "geolocation=(), microphone=(), camera=()"

    # Cross-Origin headers
    cross_origin_opener_policy: str = "same-origin"
    cross_origin_embedder_policy: str = "require-corp"
    cross_origin_resource_policy: str = "same-origin"

    # Development mode (relaxes some restrictions)
    development_mode: bool = False

    # Excluded paths (e.g., for health checks)
    excluded_paths: List[str] = field(default_factory=lambda: ["/health", "/metrics"])


class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    """
    Middleware that adds security headers to all responses.

    Attributes:
        config: SecurityHeadersConfig instance
    """

    def __init__(
        self,
        app,
        config: Optional[SecurityHeadersConfig] = None,
        # Individual overrides for convenience
        hsts_enabled: Optional[bool] = None,
        csp_policy: Optional[str] = None,
        csp_enabled: Optional[bool] = None,
        development_mode: Optional[bool] = None,
    ):
        super().__init__(app)

        # Use provided config or create default
        self.config = config or SecurityHeadersConfig()

        # Apply individual overrides
        if hsts_enabled is not None:
            self.config.hsts_enabled = hsts_enabled
        if csp_policy is not None:
            self.config.csp_policy = csp_policy
        if csp_enabled is not None:
            self.config.csp_enabled = csp_enabled
        if development_mode is not None:
            self.config.development_mode = development_mode

        # Auto-detect development mode from environment
        if development_mode is None:
            env = os.getenv("ENVIRONMENT", "development")
            self.config.development_mode = env.lower() in ("development", "dev", "local")

    def _build_hsts_header(self) -> str:
        """Build the Strict-Transport-Security header value."""
        parts = [f"max-age={self.config.hsts_max_age}"]
        if self.config.hsts_include_subdomains:
            parts.append("includeSubDomains")
        if self.config.hsts_preload:
            parts.append("preload")
        return "; ".join(parts)

    def _get_headers(self) -> Dict[str, str]:
        """Build the security headers dictionary."""
        headers = {}

        # Always add these headers
        headers["X-Content-Type-Options"] = self.config.content_type_options
        headers["X-Frame-Options"] = self.config.frame_options
        headers["X-XSS-Protection"] = self.config.xss_protection
        headers["Referrer-Policy"] = self.config.referrer_policy

        # HSTS (only in production or if explicitly enabled)
        if self.config.hsts_enabled and not self.config.development_mode:
            headers["Strict-Transport-Security"] = self._build_hsts_header()

        # Content-Security-Policy
        if self.config.csp_enabled:
            headers["Content-Security-Policy"] = self.config.csp_policy

        # Permissions-Policy
        if self.config.permissions_policy:
            headers["Permissions-Policy"] = self.config.permissions_policy

        # Cross-Origin headers (relaxed in development)
        if not self.config.development_mode:
            headers["Cross-Origin-Opener-Policy"] = self.config.cross_origin_opener_policy
            # Note: COEP can break loading of external resources, be careful
            # headers["Cross-Origin-Embedder-Policy"] = self.config.cross_origin_embedder_policy
            headers["Cross-Origin-Resource-Policy"] = self.config.cross_origin_resource_policy

        return headers

    async def dispatch(self, request: Request, call_next) -> Response:
        # Skip security headers for excluded paths
        if request.url.path in self.config.excluded_paths:
            return await call_next(request)

        # Process request
        response = await call_next(request)

        # Add security headers
        for header_name, header_value in self._get_headers().items():
            response.headers[header_name] = header_value

        return response


def get_default_csp_for_environment(environment: str) -> str:
    """
    Get a sensible default CSP for the given environment.

    Args:
        environment: "development", "staging", or "production"

    Returns:
        CSP policy string
    """
    if environment.lower() in ("development", "dev", "local"):
        # Relaxed CSP for development
        return (
            "default-src 'self' localhost:* ws://localhost:*; "
            "script-src 'self' 'unsafe-inline' 'unsafe-eval'; "
            "style-src 'self' 'unsafe-inline'; "
            "img-src 'self' data: https: blob:; "
            "font-src 'self' data:; "
            "connect-src 'self' localhost:* ws://localhost:* https:; "
            "frame-ancestors 'self'"
        )
    else:
        # Strict CSP for production
        return (
            "default-src 'self'; "
            "script-src 'self' 'unsafe-inline'; "
            "style-src 'self' 'unsafe-inline'; "
            "img-src 'self' data: https:; "
            "font-src 'self' data:; "
            "connect-src 'self' https://breakpilot.app https://*.breakpilot.app; "
            "frame-ancestors 'none'"
        )