app/middleware/csrf.py

import os
import secrets
import hashlib
from typing import Callable, Optional
from fastapi import Request, Response, HTTPException, status
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse


class CSRFMiddleware(BaseHTTPMiddleware):
    """
    CSRF Protection usando patrón Double-Submit Cookie.
    
    - Genera cookie `csrf_token` (HttpOnly=False, SameSite=Lax) al iniciar sesión
    - Valida header `X-CSRF-Token` en métodos mutantes (POST, PUT, PATCH, DELETE)
    - Excluye: GET, HEAD, OPTIONS, endpoints de auth (/api/auth/*)
    - Skip si no hay cookie (p.ej. APIs programáticas con Bearer token)
    """
    
    def __init__(
        self, 
        app, 
        cookie_name: str = "csrf_token",
        header_name: str = "X-CSRF-Token",
        cookie_secure: bool = True,
        cookie_samesite: str = "lax",
        excluded_paths: Optional[list] = None,
        excluded_methods: Optional[list] = None,
    ):
        super().__init__(app)
        self.cookie_name = cookie_name
        self.header_name = header_name
        self.cookie_secure = cookie_secure
        self.cookie_samesite = cookie_samesite
        self.excluded_paths = excluded_paths or [
            "/api/auth/jwt/login",
            "/api/auth/jwt/logout", 
            "/api/auth/register",
            "/api/auth/forgot-password",
            "/api/auth/reset-password",
            "/api/auth/users/me",
        ]
        self.excluded_methods = excluded_methods or ["GET", "HEAD", "OPTIONS"]
    
    def _get_csrf_token(self, request: Request) -> Optional[str]:
        """Obtener token CSRF de la cookie."""
        cookie_header = request.headers.get("cookie", "")
        for cookie in cookie_header.split(";"):
            cookie = cookie.strip()
            if cookie.startswith(f"{self.cookie_name}="):
                return cookie.split("=", 1)[1]
        return None
    
    def _generate_csrf_token(self) -> str:
        """Generar token CSRF criptográficamente seguro."""
        return secrets.token_urlsafe(32)
    
    def _is_excluded_path(self, path: str) -> bool:
        """Verificar si el path está excluido de validación CSRF."""
        for excluded in self.excluded_paths:
            if path.startswith(excluded):
                return True
        return False
    
    async def dispatch(self, request: Request, call_next: Callable):
        # Skip CSRF para métodos seguros
        if request.method in self.excluded_methods:
            return await call_next(request)
        
        # Skip CSRF para paths excluidos
        if self._is_excluded_path(request.url.path):
            return await call_next(request)
        
        # Solo validar si hay cookie de sesión (usuario logueado via cookie)
        # APIs programáticas con Bearer token no necesitan CSRF
        session_cookie = request.cookies.get("cd_token")
        if not session_cookie:
            # No hay cookie de sesión, asumimos API programática
            return await call_next(request)
        
        # Obtener token CSRF de cookie
        csrf_cookie = self._get_csrf_token(request)
        csrf_header = request.headers.get(self.header_name)
        
        if not csrf_cookie:
            # No hay token CSRF en cookie - generar uno nuevo y setear
            response = await call_next(request)
            new_token = self._generate_csrf_token()
            response.set_cookie(
                key=self.cookie_name,
                value=new_token,
                max_age=7 * 24 * 60 * 60,  # 7 días
                httponly=False,  # JS necesita leerlo para enviar en header
                secure=self.cookie_secure,
                samesite=self.cookie_samesite,
                path="/",
            )
            return response
        
        if not csrf_header:
            return JSONResponse(
                status_code=status.HTTP_403_FORBIDDEN,
                content={"detail": "CSRF token requerido en header X-CSRF-Token"},
            )
        
        # Validar token (comparación timing-safe)
        if not secrets.compare_digest(csrf_cookie, csrf_header):
            return JSONResponse(
                status_code=status.HTTP_403_FORBIDDEN,
                content={"detail": "CSRF token inválido"},
            )
        
        # Token válido - proceder
        return await call_next(request)


def get_csrf_middleware(
    cookie_secure: bool = True,
    cookie_samesite: str = "lax",
    excluded_paths: Optional[list] = None,
) -> type:
    """Factory para crear middleware CSRF con configuración personalizada."""
    class ConfiguredCSRFMiddleware(CSRFMiddleware):
        def __init__(self, app):
            super().__init__(
                app,
                cookie_secure=cookie_secure,
                cookie_samesite=cookie_samesite,
                excluded_paths=excluded_paths,
            )
    return ConfiguredCSRFMiddleware