orion/middleware/cloudflare.py

# middleware/cloudflare.py
"""
CloudFlare proxy middleware.

When enabled, this middleware handles CloudFlare-specific headers:
- CF-Connecting-IP: Real client IP address
- CF-IPCountry: Client's country code
- CF-Ray: CloudFlare Ray ID for request tracing
- CF-Visitor: Original request scheme

Usage:
    Enable by setting CLOUDFLARE_ENABLED=true in .env

    The middleware will:
    1. Extract real client IP from CF-Connecting-IP header
    2. Store country code in request.state.client_country
    3. Log CF-Ray for request correlation
"""

import logging

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

from app.core.config import settings

logger = logging.getLogger(__name__)


class CloudFlareMiddleware(BaseHTTPMiddleware):
    """
    Middleware to handle CloudFlare proxy headers.

    Only active when CLOUDFLARE_ENABLED=true.
    """

    async def dispatch(self, request: Request, call_next) -> Response:
        """Process request with CloudFlare headers."""

        if not settings.cloudflare_enabled:
            return await call_next(request)

        # Extract CloudFlare headers
        cf_connecting_ip = request.headers.get("CF-Connecting-IP")
        cf_country = request.headers.get("CF-IPCountry")
        cf_ray = request.headers.get("CF-Ray")
        cf_visitor = request.headers.get("CF-Visitor")  # {"scheme":"https"}

        # Store real client IP
        if cf_connecting_ip:
            # Override the client host with real IP
            # Note: Starlette's request.client is immutable, so we store in state
            request.state.real_ip = cf_connecting_ip
        else:
            # Fallback to standard client IP
            request.state.real_ip = (
                request.client.host if request.client else "unknown"
            )

        # Store country code for geo features
        if cf_country:
            request.state.client_country = cf_country

        # Store Ray ID for request tracing
        if cf_ray:
            request.state.cf_ray = cf_ray

        # Determine original scheme
        if cf_visitor:
            try:
                import json
                visitor_info = json.loads(cf_visitor)
                request.state.original_scheme = visitor_info.get("scheme", "https")
            except (json.JSONDecodeError, TypeError):
                request.state.original_scheme = "https"

        # Log request with CloudFlare context
        logger.debug(
            f"CloudFlare request: ip={cf_connecting_ip}, "
            f"country={cf_country}, ray={cf_ray}"
        )

        # Process request
        response = await call_next(request)

        # Add CF-Ray to response for debugging
        if cf_ray:
            response.headers["X-CF-Ray"] = cf_ray

        return response


def get_real_client_ip(request: Request) -> str:
    """
    Get the real client IP address.

    When behind CloudFlare, this returns CF-Connecting-IP.
    Otherwise, returns the standard client IP.

    Args:
        request: Starlette Request object

    Returns:
        Client IP address string
    """
    # First check if CloudFlare middleware set real_ip
    if hasattr(request.state, "real_ip"):
        return request.state.real_ip

    # Check for CF-Connecting-IP header directly
    cf_ip = request.headers.get("CF-Connecting-IP")
    if cf_ip:
        return cf_ip

    # Check for X-Forwarded-For (generic proxy header)
    forwarded_for = request.headers.get("X-Forwarded-For")
    if forwarded_for:
        # Take the first IP (client IP)
        return forwarded_for.split(",")[0].strip()

    # Fallback to request client
    return request.client.host if request.client else "unknown"


def get_client_country(request: Request) -> str | None:
    """
    Get the client's country code from CloudFlare.

    Args:
        request: Starlette Request object

    Returns:
        ISO 3166-1 alpha-2 country code, or None if not available
    """
    if hasattr(request.state, "client_country"):
        return request.state.client_country

    return request.headers.get("CF-IPCountry")


# =============================================================================
# PUBLIC API
# =============================================================================
__all__ = [
    "CloudFlareMiddleware",
    "get_real_client_ip",
    "get_client_country",
]