#!/usr/bin/env python3
"""HTTP server for the pageCapture user-restriction bypass POC.

Serves:
  * port 8080: the user-restricted target page (secret.html)
  * port 8081: a report endpoint the extension POSTs results to (separate
    origin so the report fetch is never blocked by the user restriction)

Writes all extension log lines to ./poc-output.log and stdout.
"""
import http.server
import os
import socketserver
import sys
import threading

ROOT = os.path.dirname(os.path.abspath(__file__))
LOG_PATH = os.path.join(ROOT, "poc-output.log")
DONE = threading.Event()


class TargetHandler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=ROOT, **kw)

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


class ReportHandler(http.server.BaseHTTPRequestHandler):
    def do_POST(self):
        length = int(self.headers.get("Content-Length", "0"))
        body = self.rfile.read(length).decode("utf-8", "replace")
        line = "[ext] " + body
        print(line, flush=True)
        with open(LOG_PATH, "a") as f:
            f.write(line + "\n")
        if "=== POC END ===" in body:
            DONE.set()
        self.send_response(200)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()

    def do_OPTIONS(self):
        self.send_response(200)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "POST")
        self.send_header("Access-Control-Allow-Headers", "*")
        self.end_headers()

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


class ReusableTCPServer(socketserver.TCPServer):
    allow_reuse_address = True


def main():
    open(LOG_PATH, "w").close()
    target = ReusableTCPServer(("127.0.0.1", 8080), TargetHandler)
    report = ReusableTCPServer(("127.0.0.1", 8081), ReportHandler)
    threading.Thread(target=target.serve_forever, daemon=True).start()
    threading.Thread(target=report.serve_forever, daemon=True).start()
    print("[server] target on http://localhost:8080/secret.html", flush=True)
    print("[server] report on http://127.0.0.1:8081/report", flush=True)
    try:
        DONE.wait()
        # Give any trailing log POSTs a moment, then exit so the runner can
        # tear down Chrome.
        threading.Event().wait(1.0)
    except KeyboardInterrupt:
        pass
    finally:
        target.shutdown()
        report.shutdown()


if __name__ == "__main__":
    main()