#!/usr/bin/env python3
"""Raw HTTP/2 (h2c prior-knowledge, or TLS via --tls) malicious-server harness.
usage: python3 h2srv.py PORT scenario.py [--tls] [--loop]
scenario.py defines run(conn, idx) where conn is an H2Conn. idx = connection number; optional hooks before_tls(idx) and use_takeover(idx).
"""
import os as _os
HERE = _os.path.dirname(_os.path.abspath(__file__))
CERT = _os.environ.get("F1_CERT", _os.path.join(HERE, "cert.pem"))
import socket, struct, threading, sys, time, importlib.util, ssl, traceback, os
from hyperframe.frame import (Frame, SettingsFrame, HeadersFrame, DataFrame, RstStreamFrame,
    GoAwayFrame, PushPromiseFrame, WindowUpdateFrame, PingFrame, PriorityFrame, ContinuationFrame)
from hpack import Encoder, Decoder

PREFACE = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"

class H2Conn:
    def __init__(self, sock, idx, verbose=True):
        self.s = sock; self.idx = idx; self.verbose = verbose
        self.enc = Encoder(); self.dec = Decoder()
        self.buf = b""; self.closed = False
        self.client_settings = {}
        self.tx = [[0, b""]]   # transcript: list of [gate, bytes]
        self.authority = None
    def log(self, *a):
        if self.verbose: print(f"[c{self.idx}]", *a, file=sys.stderr, flush=True)
    # --- raw io
    def _read(self, n, timeout=None):
        self.s.settimeout(timeout)
        while len(self.buf) < n:
            try:
                d = self.s.recv(65536)
            except socket.timeout:
                return None
            except OSError as e:
                self.log("recv err", e); self.closed = True; return b""
            if not d:
                self.closed = True; return b""
            self.buf += d
        out, self.buf = self.buf[:n], self.buf[n:]
        return out
    def read_preface(self, skip=False):
        p = self._read(len(PREFACE), timeout=10)
        if p != PREFACE:
            self.log("BAD PREFACE", p)
            return False
        self.log("preface ok")
        return True
    def recv_frame(self, timeout=5):
        hdr = self._read(9, timeout)
        if not hdr: return None
        length = struct.unpack(">I", b"\x00" + hdr[:3])[0]
        body = self._read(length, timeout)
        if body is None or (length and not body): return None
        try:
            f, l = Frame.parse_frame_header(memoryview(hdr))
            f.parse_body(memoryview(body))
        except Exception as e:
            self.log("unparseable frame", hdr.hex(), e); return ("raw", hdr, body)
        if isinstance(f, SettingsFrame) and 'ACK' not in f.flags:
            self.client_settings.update(f.settings)
        if isinstance(f, HeadersFrame):
            try:
                f.decoded = [(n.decode("latin-1"), v.decode("latin-1")) for n, v in self.dec.decode(bytes(f.data), raw=True)]
            except Exception as e:
                f.decoded = [("<decode-error>", str(e))]
            if self.tx[-1][1]: self.tx.append([1, b""])
            else: self.tx[-1][0] = 1
        self.log("<-", f, getattr(f, 'decoded', ''))
        return f
    def recv_h1_head(self, timeout=10):
        """read an HTTP/1 request head (until CRLFCRLF); returns bytes (head) ; leaves rest in buf"""
        self.s.settimeout(timeout)
        while b"\r\n\r\n" not in self.buf:
            try: d = self.s.recv(65536)
            except socket.timeout: return None
            if not d: self.closed = True; return None
            self.buf += d
        i = self.buf.index(b"\r\n\r\n") + 4
        head, self.buf = self.buf[:i], self.buf[i:]
        self.log("<- H1 head", head)
        if self.tx[-1][1]: self.tx.append([1, b""])
        else: self.tx[-1][0] = 1
        return head
    def send_raw(self, b):
        self.tx[-1][1] += bytes(b)
        try:
            self.s.sendall(b)
        except OSError as e:
            self.log("send err", e); self.closed = True
    def send(self, f):
        self.log("->", f)
        self.send_raw(f.serialize())
    # --- helpers
    def settings(self, d=None, ack=False):
        f = SettingsFrame(0)
        if ack: f.flags.add('ACK')
        else: f.settings = d or {}
        self.send(f)
    def handshake(self, settings=None):
        """read preface, send our SETTINGS, read frames until client's HEADERS arrives; ack settings. returns HeadersFrame"""
        if not self.read_preface(): return None
        self.settings(settings or {})
        return self.wait_headers()
    def wait_headers(self, timeout=10):
        while True:
            f = self.recv_frame(timeout)
            if f is None: return None
            if isinstance(f, SettingsFrame) and 'ACK' not in f.flags:
                self.settings(ack=True)
            if isinstance(f, HeadersFrame):
                return f
    def drain(self, timeout=0.3, until=None):
        """read frames for a while (auto-acking SETTINGS/PING)"""
        out = []
        end = time.time() + timeout
        while time.time() < end:
            f = self.recv_frame(max(0.01, end - time.time()))
            if f is None:
                if self.closed: break
                continue
            out.append(f)
            if isinstance(f, SettingsFrame) and 'ACK' not in f.flags:
                self.settings(ack=True)
            if isinstance(f, DataFrame) and len(f.data) and getattr(self, 'auto_wu', True):
                try:
                    self.window_update(0, len(f.data) + (f.pad_length or 0)); self.window_update(f.stream_id, len(f.data) + (f.pad_length or 0))
                except Exception: pass
            if isinstance(f, PingFrame) and 'ACK' not in f.flags:
                p = PingFrame(0); p.flags.add('ACK'); p.opaque_data = f.opaque_data; self.send(p)
            if until and until(f): break
        return out
    def hblock(self, headers, sensitive=True):
        hs = [(n if isinstance(n, bytes) else n.encode(), v if isinstance(v, bytes) else v.encode()) for n, v in headers]
        if sensitive:
            return self.enc.encode([(n, v, True) for n, v in hs])  # never-indexed literals
        return self.enc.encode(hs)
    def headers(self, sid, headers, end_stream=False, end_headers=True, priority=None, raw_block=None):
        f = HeadersFrame(sid)
        f.data = raw_block if raw_block is not None else self.hblock(headers)
        if end_stream: f.flags.add('END_STREAM')
        if end_headers: f.flags.add('END_HEADERS')
        self.send(f)
    def continuation(self, sid, block, end_headers=True):
        f = ContinuationFrame(sid); f.data = block
        if end_headers: f.flags.add('END_HEADERS')
        self.send(f)
    def data(self, sid, payload, end_stream=False, pad=None, maxf=16384):
        while len(payload) > maxf:
            f = DataFrame(sid); f.data = payload[:maxf]; payload = payload[maxf:]; self.send(f)
        f = DataFrame(sid); f.data = payload
        if pad is not None:
            f.flags.add('PADDED'); f.pad_length = pad
        if end_stream: f.flags.add('END_STREAM')
        self.send(f)
    def rst(self, sid, code):
        f = RstStreamFrame(sid); f.error_code = code; self.send(f)
    def goaway(self, last_sid, code=0, debug=b""):
        f = GoAwayFrame(0); f.last_stream_id = last_sid; f.error_code = code; f.additional_data = debug; self.send(f)
    def push(self, sid, promised, headers, end_headers=True, raw_block=None):
        f = PushPromiseFrame(sid); f.promised_stream_id = promised
        f.data = raw_block if raw_block is not None else self.hblock(headers)
        if end_headers: f.flags.add('END_HEADERS')
        self.send(f)
    def window_update(self, sid, inc):
        f = WindowUpdateFrame(sid); f.window_increment = inc; self.send(f)
    def ping(self, data=b"\0"*8):
        f = PingFrame(0); f.opaque_data = data; self.send(f)
    def priority(self, sid, dep, weight=16, excl=False):
        f = PriorityFrame(sid); f.depends_on = dep; f.stream_weight = weight; f.exclusive = excl; self.send(f)
    def rawframe(self, ftype, flags, sid, payload):
        self.log("-> raw type=%d flags=%d sid=%d len=%d" % (ftype, flags, sid, len(payload)))
        self.send_raw(struct.pack(">I", len(payload))[1:] + bytes([ftype, flags]) + struct.pack(">I", sid & 0x7fffffff) + payload)
    def resp(self, sid, status=200, hdrs=(), body=b"", end=True):
        self.headers(sid, [(":status", str(status))] + list(hdrs), end_stream=(end and not body))
        if body:
            self.data(sid, body, end_stream=end)
    def close(self):
        try: self.s.shutdown(socket.SHUT_RDWR)
        except OSError: pass
        self.s.close(); self.closed = True
        dd = os.environ.get("DUMPDIR")
        if dd and not getattr(self, "_dumped", False) and any(b for g, b in self.tx):
            self._dumped = True
            try:
                os.makedirs(dd, exist_ok=True)
                name = os.environ.get("DUMPNAME", "scen") + "-c%d" % self.idx
                with open(os.path.join(dd, name + ".chunks"), "w") as fh:
                    for g, b in self.tx:
                        if not b: continue
                        fh.write("--- conn %d gate %d\n" % (min(self.idx, 5) if self.idx < 1000 else 0, g))
                        for i in range(0, len(b), 4000):
                            fh.write(repr(b[i:i+4000]) + "\n")
            except Exception as e:
                self.log("dump failed", e)

def serve(port, scenario_path, tls=False, loop=False, host=None):
    host = host or _os.environ.get("F1_HOST", "127.0.0.1")
    spec = importlib.util.spec_from_file_location("scenario", scenario_path)
    mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod)
    ls = socket.socket(); ls.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    ls.bind((host, port)); ls.listen(int(_os.environ.get("F1_BACKLOG", "16")))
    # accept gate (the topology author, z2): the scenario can pause accepting (SERVER['pause'].set()) and park dummy self-connections
    # (SERVER['dummies'] = set of source ports) so that the kernel's accept queue is full and a client's SYN is dropped until un-paused.
    SERVER = {"pause": threading.Event(), "paused_ack": threading.Event(), "dummies": set(), "port": port, "host": host, "listener": ls}
    mod.SERVER = SERVER
    ls.settimeout(0.05)
    ctx = None
    ctxs = []
    import random as _r
    if tls:
        for al in (["h2", "http/1.1"], ["http/1.1"], ["h2"], None):
            cx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
            cx.load_cert_chain(CERT)
            if al: cx.set_alpn_protocols(al)
            cx.num_tickets = int(_os.environ.get("F1_NUM_TICKETS", "2"))
            ctxs.append(cx)
        ctx = ctxs[0]
    print(f"h2srv listening on {host}:{port} scenario={scenario_path}", file=sys.stderr, flush=True)
    idx = 0
    maxconns = getattr(mod, "MAXCONNS", 1000000 if loop else getattr(mod, "NCONNS", 1))
    threads = []
    while idx < maxconns:
        if SERVER["pause"].is_set():
            SERVER["paused_ack"].set(); time.sleep(0.01); continue   # ack: we are NOT inside accept() while paused
        SERVER["paused_ack"].clear()
        try: cs, addr = ls.accept()
        except socket.timeout: continue
        if addr[1] in SERVER["dummies"]:
            SERVER["dummies"].discard(addr[1]); cs.close(); continue   # our own queue-filler, not a client
        cs.settimeout(None)
        cs.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        if "--connect-proxy" in sys.argv:
            # act as an HTTP/1 proxy first: read CONNECT, answer 200, then continue (TLS-wrap if --tls)
            try:
                cs.settimeout(10); buf = b""
                while b"\r\n\r\n" not in buf:
                    d = cs.recv(4096)
                    if not d: break
                    buf += d
                pv = _r.random() if ctx else 0.5
                if pv < 0.1: cs.sendall(b"HTTP/1.1 100 Continue\r\n\r\nHTTP/1.1 200 OK\r\n\r\n")
                elif pv < 0.2: cs.sendall(b"HTTP/1.0 200 Connection established\r\nProxy-Agent: x\r\nContent-Length: 0\r\n\r\n")
                elif pv < 0.25: cs.sendall(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nContent-Length: 10\r\n\r\n")
                else: cs.sendall(b"HTTP/1.1 200 Connection established\r\n\r\n")
            except Exception as e:
                print("connect-proxy fail", e, file=sys.stderr); cs.close(); continue
        if "--takeover" in sys.argv or (hasattr(mod, "use_takeover") and mod.use_takeover(idx)):
            try:
                if hasattr(mod, "before_tls"): mod.before_tls(idx)
                sys.path.insert(0, HERE)
                from tls13_takeover import TLS13Takeover
                cs = TLS13Takeover(cs, CERT, alpn=("h2",)).handshake()
            except Exception as e:
                print("takeover fail", e, file=sys.stderr); 
                try: cs.close()
                except Exception: pass
                continue
        elif ctx:
            usectx = ctx
            if hasattr(mod, "before_tls"):
                try: mod.before_tls(idx)
                except Exception: traceback.print_exc()
            if "--alpn-mix" in sys.argv:
                usectx = _r.choice(ctxs if idx > 0 or _r.random() < 0.3 else ctxs[:1])
            try: cs = usectx.wrap_socket(cs, server_side=True)
            except Exception as e: print("tls fail", e, file=sys.stderr); cs.close(); continue
        conn = H2Conn(cs, idx, verbose=getattr(mod, "VERBOSE", True))
        def runit(conn=conn, i=idx):
            try: mod.run(conn, i)
            except Exception: traceback.print_exc()
            finally:
                try: conn.close()
                except Exception: pass
        t = threading.Thread(target=runit, daemon=True); t.start(); threads.append(t)
        idx += 1
    for t in threads: t.join(30)

if __name__ == "__main__":
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    serve(int(args[0]), args[1], tls="--tls" in sys.argv, loop="--loop" in sys.argv)