#!/usr/bin/env python3
"""
Generate netbios-truncated-session-msg.pcap: a NetBIOS-over-TCP session in
which the client declares the maximum possible message size (0xFFFF = 65535
bytes) in the four-byte session header but then sends only a partial body
before closing the connection.

The four-byte NetBIOS session header is: type (1 byte), flags (1 byte),
length (2 bytes, big-endian). The length field is 16 bits, so 0xFFFF is the
largest message size it can express.

This exercises Contents_NetbiosSSN::Flush(): on connection close the analyzer
delivers the partial buffer. With buf_n < msg_size it should emit the
truncated_netbios_message weird. A fully-satisfied message would not.

Generated by Claude Opus 4.8 (Anthropic). Based on a GPT-5.6 sol generated
script from Trail of Bits.
"""

import argparse
from pathlib import Path

from scapy.all import IP, TCP, Ether, wrpcap

CLIENT_IP = "10.0.0.1"
SERVER_IP = "10.0.0.2"
CLIENT_PORT = 54321
SERVER_PORT = 139  # NetBIOS session service

# NetBIOS session message type; 16-bit length field maxes out at 0xFFFF.
NETBIOS_SSN_MESSAGE = 0x00
MAX_MSG_SIZE = 0xFFFF


def pkt(src_ip, dst_ip, sport, dport, flags, seq, ack, payload=b""):
    return (
        Ether(src="00:00:00:00:00:01", dst="00:00:00:00:00:02")
        / IP(src=src_ip, dst=dst_ip)
        / TCP(sport=sport, dport=dport, flags=flags, seq=seq, ack=ack)
        / payload
    )


def cpkt(flags, seq, ack, payload=b""):
    return pkt(CLIENT_IP, SERVER_IP, CLIENT_PORT, SERVER_PORT, flags, seq, ack, payload)


def spkt(flags, seq, ack, payload=b""):
    return pkt(SERVER_IP, CLIENT_IP, SERVER_PORT, CLIENT_PORT, flags, seq, ack, payload)


def netbios_header(msg_size):
    """Four-byte NetBIOS session header declaring msg_size body bytes."""
    return bytes([NETBIOS_SSN_MESSAGE, 0x00, (msg_size >> 8) & 0xFF, msg_size & 0xFF])


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output",
        type=Path,
        default=Path(__file__).with_suffix(""),
        help="output pcap path (default: this script's name without .py)",
    )
    parser.add_argument(
        "--msg-size",
        type=int,
        default=MAX_MSG_SIZE,
        help=f"declared NetBIOS message size (default: {MAX_MSG_SIZE}, the 16-bit max)",
    )
    parser.add_argument(
        "--body-len",
        type=int,
        default=100,
        help="actual body bytes to send before closing (default: 100)",
    )
    args = parser.parse_args()

    if not 0 <= args.msg_size <= MAX_MSG_SIZE:
        parser.error(f"--msg-size must be between 0 and {MAX_MSG_SIZE}")
    if args.body_len < 0:
        parser.error("--body-len must be non-negative")

    # Declare msg_size in the header, then send only body_len body bytes.
    payload = netbios_header(args.msg_size) + b"A" * args.body_len

    c_seq = 1000
    s_seq = 2000

    pkts = []

    # Three-way handshake so the stream reassembler engages.
    pkts.append(cpkt("S", c_seq, 0))
    c_seq += 1
    pkts.append(spkt("SA", s_seq, c_seq))
    s_seq += 1
    pkts.append(cpkt("A", c_seq, s_seq))

    # Client sends header + partial body.
    pkts.append(cpkt("PA", c_seq, s_seq, payload))
    c_seq += len(payload)
    pkts.append(spkt("A", s_seq, c_seq))

    # Teardown — triggers EndpointEOF -> Flush() with buf_n < msg_size.
    pkts.append(cpkt("FA", c_seq, s_seq))
    c_seq += 1
    pkts.append(spkt("FA", s_seq, c_seq))
    s_seq += 1
    pkts.append(cpkt("A", c_seq, s_seq))

    wrpcap(str(args.output), pkts)
    print(f"Wrote {len(pkts)} packets to {args.output}")


if __name__ == "__main__":
    main()
