#!/usr/bin/env python3
"""Focused stress test for the IPv6 host-counters map race: floods distinct IPv6 source addresses via NetFlow v9
to force continuous insertions into ipv6_host_counters.average_speed_map, racing against
the periodic (check_period=1s) unlocked iteration in draw_table_ipv6()."""
import socket, struct, sys, threading, time, random

DEST_IP = "127.0.0.1"
PORT = 2055
SOURCE_ID = 42
TEMPLATE_ID = 700
stop_flag = threading.Event()


def netflow_header(flowset_count, seq):
    return struct.pack(">HHIIII", 9, flowset_count, 0, 0, seq, SOURCE_ID)


def ipv6_template_flowset(template_id):
    body = struct.pack(">HH", template_id, 1) + struct.pack(">HH", 27, 16)  # NETFLOW9_IPV6_SRC_ADDR
    return struct.pack(">HH", 0, 4 + len(body)) + body


def ipv6_data_flowset(template_id, n_records=16):
    body = b""
    for _ in range(n_records):
        body += bytes([0x20, 0x01, 0x0d, 0xb8]) + bytes(random.randint(0, 255) for _ in range(12))
    return struct.pack(">HH", template_id, 4 + len(body)) + body


def worker(sock, duration, thread_id):
    seq = thread_id * 1000000
    count = 0
    end = time.time() + duration
    while time.time() < end and not stop_flag.is_set():
        seq += 1
        if count % 50 == 0:
            pkt = netflow_header(1, seq) + ipv6_template_flowset(TEMPLATE_ID)
            try: sock.sendto(pkt, (DEST_IP, PORT))
            except OSError: pass
        pkt2 = netflow_header(1, seq) + ipv6_data_flowset(TEMPLATE_ID)
        try: sock.sendto(pkt2, (DEST_IP, PORT))
        except OSError: pass
        count += 1
    print(f"[worker-{thread_id}] sent {count} data packets")


def main():
    duration = float(sys.argv[1]) if len(sys.argv) > 1 else 60
    threads = [threading.Thread(target=worker, args=(socket.socket(socket.AF_INET, socket.SOCK_DGRAM), duration, i))
               for i in range(4)]
    print(f"Starting 4 IPv6-flood threads for {duration}s targeting {DEST_IP}:{PORT}...")
    start = time.time()
    for t in threads: t.start()
    for t in threads: t.join()
    print(f"Done in {time.time() - start:.1f}s")


if __name__ == "__main__":
    main()