#!/usr/bin/env python3
"""
Raw TCP / TLS client churn against the target node's P2P port.

dragonxd P2P is TLS 1.3 only (wolfSSL), self-signed, no cert verification. The
accept path is net.cpp CreateNodeFromAcceptedSocket -> tlsmanager.accept()
(wolfSSL_accept), and teardown is CloseSocketDisconnect / ~CNode
(wolfSSL_shutdown + wolfSSL_free + CloseSocket). This driver pokes the accept and
TLS-teardown/error paths (SSL_read/SSL_write at net.cpp ~1005 and tlsmanager.cpp
~660-680) by opening many short-lived connections that each misbehave:

  rst      : connect, then abort with RST (SO_LINGER 0) before any TLS bytes
  fin      : connect, then close cleanly (FIN) before any TLS bytes
  garbage  : connect, send non-ClientHello bytes, then drop
  stall    : connect, send a few bytes, hold mid-handshake, drop
  tlsjunk  : complete a real TLS 1.3 handshake (verify off), send junk, drop
  tlshalf  : start TLS, send a partial ClientHello record, drop

Loopback only, against a port this harness started. --duration seconds.
"""
import argparse, random, socket, ssl, struct, sys, time, os

def mk():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(3.0)
    return s

def rst_close(s):
    try:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
    except Exception:
        pass
    try: s.close()
    except Exception: pass

def do_rst(host, port):
    s = mk()
    try:
        s.connect((host, port)); rst_close(s)
    except Exception:
        rst_close(s)

def do_fin(host, port):
    s = mk()
    try:
        s.connect((host, port)); s.close()
    except Exception:
        try: s.close()
        except Exception: pass

def do_garbage(host, port):
    s = mk()
    try:
        s.connect((host, port))
        s.sendall(os.urandom(random.randint(16, 400)))
        time.sleep(random.uniform(0.0, 0.3))
    except Exception:
        pass
    rst_close(s)

def do_stall(host, port):
    s = mk()
    try:
        s.connect((host, port))
        # a couple of bytes that look like the start of a TLS record, then wait
        s.sendall(b"\x16\x03\x01")
        time.sleep(random.uniform(0.5, 2.0))
    except Exception:
        pass
    rst_close(s)

def _ctx():
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    try:
        ctx.minimum_version = ssl.TLSVersion.TLSv1_3
    except Exception:
        pass
    return ctx

def do_tlsjunk(host, port):
    raw = mk()
    try:
        raw.connect((host, port))
        tls = _ctx().wrap_socket(raw, server_hostname=None)
        # handshake completed; the node now expects a P2P version message.
        # send junk inside the TLS channel, then tear down.
        try:
            tls.sendall(os.urandom(random.randint(8, 128)))
        except Exception:
            pass
        time.sleep(random.uniform(0.0, 0.4))
        try:
            tls.close()
        except Exception:
            rst_close(raw)
    except Exception:
        rst_close(raw)

def do_tlshalf(host, port):
    s = mk()
    try:
        s.connect((host, port))
        # a well-formed TLS record header claiming a long handshake, but we send
        # only part of it, then drop -> server sits in wolfSSL_accept WANT_READ.
        body = os.urandom(40)
        rec = b"\x16\x03\x03" + struct.pack(">H", 512) + body
        s.sendall(rec)
        time.sleep(random.uniform(0.3, 1.5))
    except Exception:
        pass
    rst_close(s)

HANDLERS = [do_rst, do_fin, do_garbage, do_stall, do_tlsjunk, do_tlshalf]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, required=True)
    ap.add_argument("--duration", type=float, default=300.0)
    ap.add_argument("--concurrency", type=int, default=6)
    ap.add_argument("--rate", type=float, default=0.03, help="sleep between bursts")
    args = ap.parse_args()

    end = time.time() + args.duration
    counts = {h.__name__: 0 for h in HANDLERS}
    import threading
    lock = threading.Lock()

    def worker():
        while time.time() < end:
            h = random.choice(HANDLERS)
            h(args.host, args.port)
            with lock:
                counts[h.__name__] += 1
            time.sleep(random.uniform(0.0, args.rate))

    ts = [threading.Thread(target=worker, daemon=True) for _ in range(args.concurrency)]
    for t in ts: t.start()
    for t in ts: t.join()

    total = sum(counts.values())
    print(f"tcpchurn: {total} connections in {args.duration:.0f}s -> {counts}", flush=True)

if __name__ == "__main__":
    main()