#!/usr/bin/env python3
"""
Dual-origin HTTP server for the TextSuggestionHostAndroid confused-deputy POC.

Serves:
  - attacker.html  on  http://<host>:8080/attacker.html   (attacker origin)
  - victim.html    on  http://<host>:8081/victim.html     (victim  origin)
  - /gen/**        on  both ports, mapped to <chromium>/out/<dir>/gen/**
                   so the page can import the generated MojoJS bindings
                   (mojo/public/js/bindings.js and *.mojom.m.js).

The two ports give two distinct web origins, so with --site-per-process the
victim iframe is rendered in a separate renderer process (true OOPIF).

Usage:
  python3 server.py [--gen-dir /path/to/chromium/src/out/asan/gen]

When testing on a physical Android device, use `adb reverse` to expose both
ports on the device's localhost (see README.md).
"""

import argparse
import http.server
import os
import socketserver
import threading

DEFAULT_GEN_DIRS = [
    "/chromium/chromium/src/out/asan/gen",
    "/chromium/chromium/src/out/android_asan/gen",
    "/chromium/chromium/src/out/Default/gen",
]

ARGS = argparse.Namespace(gen_dir=None)
HERE = os.path.dirname(os.path.abspath(__file__))


def resolve_gen_dir(explicit):
    if explicit:
        return explicit
    for d in DEFAULT_GEN_DIRS:
        if os.path.isfile(os.path.join(d, "mojo/public/js/bindings.js")):
            return d
    raise SystemExit(
        "Could not locate a gen/ dir containing mojo/public/js/bindings.js. "
        "Pass --gen-dir explicitly."
    )


class Handler(http.server.SimpleHTTPRequestHandler):
    # Filled in by make_handler().
    victim_origin = None

    def translate_path(self, path):
        # Strip query/fragment.
        path = path.split("?", 1)[0].split("#", 1)[0]
        if path.startswith("/gen/"):
            rel = path[len("/gen/"):]
            return os.path.join(ARGS.gen_dir, rel)
        if path == "/" or path == "":
            path = "/attacker.html"
        return os.path.join(HERE, path.lstrip("/"))

    def do_GET(self):
        if self.path.split("?", 1)[0] in ("/", "/attacker.html"):
            return self._serve_attacker()
        return super().do_GET()

    def _serve_attacker(self):
        with open(os.path.join(HERE, "attacker.html"), "rb") as f:
            body = f.read()
        body = body.replace(b"VICTIM_ORIGIN_PLACEHOLDER",
                            self.victim_origin.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)

    def end_headers(self):
        # Allow cross-origin embedding of victim.html.
        self.send_header("Access-Control-Allow-Origin", "*")
        super().end_headers()

    def guess_type(self, path):
        if path.endswith(".js") or path.endswith(".m.js"):
            return "text/javascript"
        return super().guess_type(path)


class ReusableTCPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True


def make_handler(victim_origin):
    return type("BoundHandler", (Handler,), {"victim_origin": victim_origin})


def serve(port, victim_origin):
    httpd = ReusableTCPServer(("127.0.0.1", port), make_handler(victim_origin))
    print(f"[*] listening on http://127.0.0.1:{port}/")
    httpd.serve_forever()


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--gen-dir", default=None,
                   help="Path to <chromium>/out/<builddir>/gen")
    p.add_argument("--attacker-port", type=int, default=8080)
    p.add_argument("--victim-port", type=int, default=8081)
    p.add_argument("--victim-host", default="127.0.0.1",
                   help="Host the attacker page uses to load the victim "
                        "iframe (use the device-visible host when testing "
                        "on Android without adb reverse).")
    global ARGS
    ARGS = p.parse_args()
    ARGS.gen_dir = resolve_gen_dir(ARGS.gen_dir)
    print(f"[*] gen dir: {ARGS.gen_dir}")

    victim_origin = f"{ARGS.victim_host}:{ARGS.victim_port}"

    t = threading.Thread(
        target=serve, args=(ARGS.victim_port, victim_origin), daemon=True
    )
    t.start()
    serve(ARGS.attacker_port, victim_origin)


if __name__ == "__main__":
    main()