proxy.py

"""HTML proxy endpoint for interactive plots with size reporting."""

import json
from urllib.parse import urlparse

import httpx
from fastapi import APIRouter, HTTPException
from fastapi.responses import HTMLResponse


router = APIRouter(tags=["proxy"])

# Allowed origins for postMessage
ALLOWED_ORIGINS = ["https://pyplots.ai", "http://localhost:3000"]


def get_size_reporter_script(target_origin: str) -> str:
    """Generate size reporter script with specified target origin.

    Uses json.dumps() to safely encode the origin string for JavaScript context,
    preventing XSS even if the origin contained special characters.
    """
    # Safely encode for JavaScript string context (includes quotes)
    safe_origin = json.dumps(target_origin)
    return f"""
<script>
(function() {{
  function reportSize() {{
    try {{
      // Find the main content element (try common patterns for different libraries)
      var content = document.querySelector(
        '.bk-root, .vega-embed, .plotly, .chart-container, #container, .lp-plot, svg, canvas'
      ) || document.body.firstElementChild || document.body;

      // Get actual rendered size
      var rect = content.getBoundingClientRect();
      var width = Math.max(rect.width, content.scrollWidth || 0, document.body.scrollWidth || 0);
      var height = Math.max(rect.height, content.scrollHeight || 0, document.body.scrollHeight || 0);

      // Add padding to account for action buttons, toolbars, and other UI elements
      var padding = 40;
      width += padding;
      height += padding;

      // Send to parent with specific origin for security
      if (width > 0 && height > 0 && window.parent !== window) {{
        window.parent.postMessage({{
          type: 'pyplots-size',
          width: Math.ceil(width),
          height: Math.ceil(height)
        }}, {safe_origin});
      }}
    }} catch (e) {{
      // Silently fail if postMessage is blocked
    }}
  }}

  // Report after load and after delays (for async rendering libraries)
  if (document.readyState === 'complete') {{
    setTimeout(reportSize, 100);
    setTimeout(reportSize, 500);
    setTimeout(reportSize, 1000);
  }} else {{
    window.addEventListener('load', function() {{
      setTimeout(reportSize, 100);
      setTimeout(reportSize, 500);
      setTimeout(reportSize, 1000);
    }});
  }}
}})();
</script>
"""


# Legacy constant for backwards compatibility with tests
SIZE_REPORTER_SCRIPT = get_size_reporter_script("https://pyplots.ai")

# Allowed GCS bucket for security
ALLOWED_HOST = "storage.googleapis.com"
ALLOWED_BUCKET = "pyplots-images"


def build_safe_gcs_url(url: str) -> str | None:
    """
    Validate URL and return a reconstructed safe GCS URL.

    This prevents SSRF by constructing the URL from hardcoded values
    instead of passing user input directly.

    Args:
        url: User-provided URL to validate

    Returns:
        Reconstructed safe URL or None if validation fails
    """
    try:
        parsed = urlparse(url)
        # Must be HTTPS
        if parsed.scheme != "https":
            return None
        # Must be exact host (no subdomains)
        if parsed.netloc != ALLOWED_HOST:
            return None
        # Path must start with bucket name
        path_parts = parsed.path.strip("/").split("/")
        if len(path_parts) < 2:
            return None
        if path_parts[0] != ALLOWED_BUCKET:
            return None
        # Check for path traversal attempts
        if ".." in parsed.path:
            return None
        # Validate path contains only safe characters (alphanumeric, hyphens, underscores, dots, slashes)
        safe_path = parsed.path.strip("/")
        if not all(c.isalnum() or c in "-_./+" for c in safe_path):
            return None
        # Reconstruct URL from hardcoded values to prevent SSRF
        # This breaks the taint flow by not using the original URL
        return f"https://{ALLOWED_HOST}/{safe_path}"
    except Exception:
        return None


@router.get("/proxy/html", response_class=HTMLResponse)
async def proxy_html(url: str, origin: str | None = None):
    """
    Proxy an HTML file and inject size reporting script.

    This endpoint fetches HTML from GCS, injects a script that reports
    the content's actual dimensions via postMessage, and returns the
    modified HTML. This allows the frontend to dynamically scale the
    iframe based on actual content size.

    Args:
        url: The GCS URL to fetch (must be from allowed bucket)
        origin: Target origin for postMessage (must be in ALLOWED_ORIGINS)

    Returns:
        Modified HTML with size reporting script injected
    """
    # Security: Validate and reconstruct URL to prevent SSRF
    safe_url = build_safe_gcs_url(url)
    if safe_url is None:
        raise HTTPException(status_code=400, detail=f"Only URLs from {ALLOWED_HOST}/{ALLOWED_BUCKET} are allowed")

    # Validate origin parameter - default to production if not specified or invalid
    target_origin = "https://pyplots.ai"
    if origin and origin in ALLOWED_ORIGINS:
        target_origin = origin

    # Fetch the HTML with shorter timeout
    async with httpx.AsyncClient(timeout=10.0) as client:
        try:
            response = await client.get(safe_url)
            response.raise_for_status()
        except httpx.HTTPStatusError as e:
            raise HTTPException(status_code=e.response.status_code, detail="Failed to fetch HTML") from e
        except httpx.RequestError as e:
            raise HTTPException(status_code=502, detail="Failed to connect to storage") from e

    html_content = response.text

    # Generate script with correct target origin
    size_script = get_size_reporter_script(target_origin)

    # Inject the size reporter script before </body>
    if "</body>" in html_content:
        html_content = html_content.replace("</body>", f"{size_script}</body>")
    elif "</html>" in html_content:
        html_content = html_content.replace("</html>", f"{size_script}</html>")
    else:
        # Fallback: append to end
        html_content += size_script

    return HTMLResponse(content=html_content)