Timmy-time-dashboard/src/dashboard/middleware/csrf.py

"""CSRF protection middleware for FastAPI.

Provides CSRF token generation, validation, and middleware integration
to protect state-changing endpoints from cross-site request attacks.
"""

import secrets
import hmac
import hashlib
from typing import Callable, Optional
from functools import wraps

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


# Module-level set to track exempt routes
_exempt_routes: set[str] = set()


def csrf_exempt(endpoint: Callable) -> Callable:
    """Decorator to mark an endpoint as exempt from CSRF validation.
    
    Usage:
        @app.post("/webhook")
        @csrf_exempt
        def webhook_endpoint():
            ...
    """
    @wraps(endpoint)
    async def async_wrapper(*args, **kwargs):
        return await endpoint(*args, **kwargs)
    
    @wraps(endpoint)
    def sync_wrapper(*args, **kwargs):
        return endpoint(*args, **kwargs)
    
    # Mark the original function as exempt
    endpoint._csrf_exempt = True  # type: ignore
    
    # Also mark the wrapper
    if hasattr(endpoint, '__code__') and endpoint.__code__.co_flags & 0x80:
        async_wrapper._csrf_exempt = True  # type: ignore
        return async_wrapper
    else:
        sync_wrapper._csrf_exempt = True  # type: ignore
        return sync_wrapper


def is_csrf_exempt(endpoint: Callable) -> bool:
    """Check if an endpoint is marked as CSRF exempt."""
    return getattr(endpoint, '_csrf_exempt', False)


def generate_csrf_token() -> str:
    """Generate a cryptographically secure CSRF token.
    
    Returns:
        A secure random token string.
    """
    return secrets.token_urlsafe(32)


def validate_csrf_token(token: str, expected_token: str) -> bool:
    """Validate a CSRF token against the expected token.
    
    Uses constant-time comparison to prevent timing attacks.
    
    Args:
        token: The token provided by the client.
        expected_token: The expected token (from cookie/session).
        
    Returns:
        True if the token is valid, False otherwise.
    """
    if not token or not expected_token:
        return False
    
    return hmac.compare_digest(token, expected_token)


class CSRFMiddleware(BaseHTTPMiddleware):
    """Middleware to enforce CSRF protection on state-changing requests.
    
    Safe methods (GET, HEAD, OPTIONS, TRACE) are allowed without CSRF tokens.
    State-changing methods (POST, PUT, DELETE, PATCH) require a valid CSRF token.
    
    The token is expected to be:
    - In the X-CSRF-Token header, or
    - In the request body as 'csrf_token', or
    - Matching the token in the csrf_token cookie
    
    Usage:
        app.add_middleware(CSRFMiddleware, secret="your-secret-key")
    
    Attributes:
        secret: Secret key for token signing (optional, for future use).
        cookie_name: Name of the CSRF cookie.
        header_name: Name of the CSRF header.
        safe_methods: HTTP methods that don't require CSRF tokens.
    """
    
    SAFE_METHODS = {"GET", "HEAD", "OPTIONS", "TRACE"}
    
    def __init__(
        self,
        app,
        secret: Optional[str] = None,
        cookie_name: str = "csrf_token",
        header_name: str = "X-CSRF-Token",
        form_field: str = "csrf_token"
    ):
        super().__init__(app)
        self.secret = secret
        self.cookie_name = cookie_name
        self.header_name = header_name
        self.form_field = form_field
    
    async def dispatch(self, request: Request, call_next) -> Response:
        """Process the request and enforce CSRF protection.
        
        For safe methods: Set a CSRF token cookie if not present.
        For unsafe methods: Validate the CSRF token.
        """
        # Bypass CSRF if explicitly disabled (e.g. in tests)
        import os
        if os.environ.get("TIMMY_DISABLE_CSRF") == "1":
            return await call_next(request)

        # Get existing CSRF token from cookie
        csrf_cookie = request.cookies.get(self.cookie_name)
        
        # For safe methods, just ensure a token exists
        if request.method in self.SAFE_METHODS:
            response = await call_next(request)
            
            # Set CSRF token cookie if not present
            if not csrf_cookie:
                new_token = generate_csrf_token()
                response.set_cookie(
                    key=self.cookie_name,
                    value=new_token,
                    httponly=False,  # Must be readable by JavaScript
                    secure=False,    # Set to True in production with HTTPS
                    samesite="Lax",
                    max_age=86400    # 24 hours
                )
            
            return response
        
        # For unsafe methods, check if route is exempt first
        # Note: We need to let the request proceed and check at response time
        # since FastAPI routes are resolved after middleware
        
        # Try to validate token early
        if not await self._validate_request(request, csrf_cookie):
            # Check if this might be an exempt route by checking path patterns
            # that are commonly exempt (like webhooks)
            path = request.url.path
            if not self._is_likely_exempt(path):
                return JSONResponse(
                    status_code=403,
                    content={
                        "error": "CSRF validation failed",
                        "code": "CSRF_INVALID",
                        "message": "Missing or invalid CSRF token. Include the token from the csrf_token cookie in the X-CSRF-Token header or as a form field."
                    }
                )
        
        return await call_next(request)
    
    def _is_likely_exempt(self, path: str) -> bool:
        """Check if a path is likely to be CSRF exempt.
        
        Common patterns like webhooks, API endpoints, etc.
        
        Args:
            path: The request path.
            
        Returns:
            True if the path is likely exempt.
        """
        exempt_patterns = [
            "/webhook",
            "/api/v1/",
            "/lightning/webhook",
            "/_internal/",
        ]
        return any(path.startswith(pattern) for pattern in exempt_patterns)
    
    async def _validate_request(self, request: Request, csrf_cookie: Optional[str]) -> bool:
        """Validate the CSRF token in the request.
        
        Checks for token in:
        1. X-CSRF-Token header
        2. csrf_token form field
        
        Args:
            request: The incoming request.
            csrf_cookie: The expected token from the cookie.
            
        Returns:
            True if the token is valid, False otherwise.
        """
        # Validate against cookie
        if not csrf_cookie:
            return False

        # Get token from header
        header_token = request.headers.get(self.header_name)
        if header_token and validate_csrf_token(header_token, csrf_cookie):
            return True
        
        # If no header token, try form data (for non-JSON POSTs)
        # Check Content-Type to avoid hanging on non-form requests
        content_type = request.headers.get("Content-Type", "")
        if "application/x-www-form-urlencoded" in content_type or "multipart/form-data" in content_type:
            try:
                form_data = await request.form()
                form_token = form_data.get(self.form_field)
                if form_token and validate_csrf_token(str(form_token), csrf_cookie):
                    return True
            except Exception:
                # Error parsing form data, treat as invalid
                pass
        
        return False