platform/gateway/integrations/image_proxy.py

"""
Platform Gateway — Image proxy with domain allowlist.
"""

import urllib.request
import urllib.parse

from config import _ssl_ctx

ALLOWED_IMAGE_DOMAINS = {
    "i.redd.it",
    "preview.redd.it",
    "external-preview.redd.it",
    "i.imgur.com",
    "imgur.com",
    "images-na.ssl-images-amazon.com",
    "m.media-amazon.com",
    "cdn.discordapp.com",
    "media.discordapp.net",
    "pbs.twimg.com",
    "abs.twimg.com",
    "upload.wikimedia.org",
    "images.unsplash.com",
}


def handle_image_proxy(handler):
    """Proxy external images to bypass hotlink protection."""
    qs = urllib.parse.urlparse(handler.path).query
    params = urllib.parse.parse_qs(qs)
    url = params.get("url", [None])[0]
    if not url:
        handler._send_json({"error": "Missing url parameter"}, 400)
        return

    # Validate URL scheme and domain allowlist
    try:
        parsed = urllib.parse.urlparse(url)
        if parsed.scheme not in ("http", "https"):
            handler._send_json({"error": "Invalid URL scheme"}, 400)
            return
        hostname = parsed.hostname or ""
        if hostname not in ALLOWED_IMAGE_DOMAINS:
            allowed = any(
                hostname == d or hostname.endswith(f".{d}")
                for d in ALLOWED_IMAGE_DOMAINS
            )
            if not allowed:
                print(f"[ImageProxy] Blocked request for domain: {hostname}")
                handler._send_json({"error": "Domain not allowed"}, 403)
                return
    except Exception:
        handler._send_json({"error": "Invalid URL"}, 400)
        return

    try:
        req = urllib.request.Request(url, headers={
            "User-Agent": "Mozilla/5.0 (compatible; PlatformProxy/1.0)",
            "Accept": "image/*,*/*",
            "Referer": parsed.scheme + "://" + parsed.netloc + "/",
        })
        resp = urllib.request.urlopen(req, timeout=10, context=_ssl_ctx)
        body = resp.read()
        ct = resp.headers.get("Content-Type", "image/jpeg")
        # Only serve actual image content types
        if not ct.startswith("image/"):
            handler._send_json({"error": "Response is not an image"}, 400)
            return
        handler.send_response(200)
        handler.send_header("Content-Type", ct)
        handler.send_header("Content-Length", len(body))
        handler.send_header("Cache-Control", "public, max-age=86400")
        handler.end_headers()
        handler.wfile.write(body)
    except Exception as e:
        print(f"[ImageProxy] Error fetching {url}: {e}")
        handler._send_json({"error": "Failed to fetch image"}, 502)