#!/usr/bin/env python3
"""
HTTP server for the MIME map race UAF POC.

Serves poc.html plus the generated MojoJS bindings from the Chromium
build output directory. Listens on 127.0.0.1:8080.
"""
import http.server
import os
import socketserver
import sys

PORT = 18080
POC_DIR = os.path.dirname(os.path.abspath(__file__))

# Generated MojoJS bindings live under the build output's gen/ directory.
# Override via CHROMIUM_OUT_DIR if your build is elsewhere.
GEN_DIR = os.environ.get(
    "CHROMIUM_OUT_DIR",
    "/chromium/chromium/src/out/asan",
) + "/gen"


class Handler(http.server.SimpleHTTPRequestHandler):
    def translate_path(self, path):
        # Route /gen/* to the Chromium build's generated-files directory so
        # poc.html can load mojo_bindings_lite.js + mime_registry.mojom-lite.js.
        if path.startswith("/gen/"):
            rel = path[len("/gen/"):]
            # Disallow path escapes.
            full = os.path.normpath(os.path.join(GEN_DIR, rel))
            if not full.startswith(os.path.normpath(GEN_DIR)):
                return "/dev/null"
            return full
        # Everything else from the POC directory.
        return os.path.join(POC_DIR, path.lstrip("/"))

    def end_headers(self):
        # MojoJS bindings load fine without special headers, but set a
        # permissive CSP just in case.
        self.send_header("Cache-Control", "no-store")
        super().end_headers()

    def log_message(self, fmt, *args):
        sys.stderr.write("[server] %s\n" % (fmt % args))


if __name__ == "__main__":
    socketserver.TCPServer.allow_reuse_address = True
    with socketserver.TCPServer(("127.0.0.1", PORT), Handler) as httpd:
        sys.stderr.write(f"[server] listening on http://127.0.0.1:{PORT}\n")
        sys.stderr.write(f"[server] gen dir: {GEN_DIR}\n")
        httpd.serve_forever()