#!/usr/bin/env python3
from __future__ import annotations

import argparse
import socket
import struct
import time


def checksum(data: bytes) -> int:
    if len(data) % 2:
        data += b"\x00"
    total = sum(struct.unpack(f"!{len(data) // 2}H", data))
    total = (total >> 16) + (total & 0xFFFF)
    total += total >> 16
    return (~total) & 0xFFFF


def ipv4_bytes(address: str) -> bytes:
    return socket.inet_aton(address)


def build_icmp_frame(src_ip: str, dst_ip: str, sequence: int) -> bytes:
    dst_mac = bytes.fromhex("020000000002")
    src_mac = bytes.fromhex("020000000001")
    ethernet = dst_mac + src_mac + struct.pack("!H", 0x0800)

    payload = b"routeros-suricata-tzsp-selftest"
    icmp_header = struct.pack("!BBHHH", 8, 0, 0, 0x1234, sequence & 0xFFFF)
    icmp_sum = checksum(icmp_header + payload)
    icmp = struct.pack("!BBHHH", 8, 0, icmp_sum, 0x1234, sequence & 0xFFFF) + payload

    total_length = 20 + len(icmp)
    ip_header = struct.pack(
        "!BBHHHBBH4s4s",
        0x45,
        0,
        total_length,
        sequence & 0xFFFF,
        0,
        64,
        socket.IPPROTO_ICMP,
        0,
        ipv4_bytes(src_ip),
        ipv4_bytes(dst_ip),
    )
    ip_sum = checksum(ip_header)
    ip_header = struct.pack(
        "!BBHHHBBH4s4s",
        0x45,
        0,
        total_length,
        sequence & 0xFFFF,
        0,
        64,
        socket.IPPROTO_ICMP,
        ip_sum,
        ipv4_bytes(src_ip),
        ipv4_bytes(dst_ip),
    )
    return ethernet + ip_header + icmp


def wrap_tzsp(frame: bytes) -> bytes:
    # Version=1, Type=0 (received packet), Encapsulation=1 (Ethernet), END tag=1.
    return b"\x01\x00\x00\x01\x01" + frame


def main() -> int:
    parser = argparse.ArgumentParser(description="Send deterministic TZSP Ethernet frames for Suricata testing")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=37008)
    parser.add_argument("--count", type=int, default=3)
    parser.add_argument("--src-ip", default="192.168.100.10")
    parser.add_argument("--dst-ip", default="1.1.1.1")
    args = parser.parse_args()

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    for seq in range(1, args.count + 1):
        packet = wrap_tzsp(build_icmp_frame(args.src_ip, args.dst_ip, seq))
        sock.sendto(packet, (args.host, args.port))
        print(f"sent TZSP test datagram {seq}/{args.count} to {args.host}:{args.port}")
        time.sleep(0.1)
    sock.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())