#!/usr/bin/env python3
"""Compute record-boundary-aligned byte offsets in a classic (non-pcapng)
.pcap file, for simulating a still-growing capture in regression tests.

Usage: pcap_offsets.py <file> <n_segments>

Prints n_segments+1 offsets, one per line: 0 (start of body, i.e. just past
the 24-byte global header) ... file size (end), each guaranteed to land
exactly on a record boundary. Segment sizes are as even as possible by
record count. Exits nonzero with a message on stderr if the file isn't a
well-formed classic pcap (e.g. truncated mid-record, or pcapng).
"""
import struct
import sys

MAGICS = {
    0xA1B2C3D4: ("<", False),  # little-endian, microsecond
    0xA1B23C4D: ("<", False),  # little-endian, nanosecond (same header shape)
    0xD4C3B2A1: (">", False),  # big-endian, microsecond
    0x4D3CB2A1: (">", False),  # big-endian, nanosecond
}


def record_offsets(path):
    with open(path, "rb") as f:
        data = f.read()
    if len(data) < 24:
        sys.exit(f"file too short to contain a pcap global header: {path}")
    magic = struct.unpack("<I", data[0:4])[0]
    if magic not in MAGICS:
        sys.exit(f"unrecognized pcap magic {magic:#010x} in {path} (pcapng is not supported)")
    endian, _ = MAGICS[magic]
    offsets = [24]
    pos = 24
    while pos < len(data):
        if pos + 16 > len(data):
            sys.exit(f"file ends mid-record-header at offset {pos}: {path}")
        incl_len = struct.unpack(endian + "I", data[pos + 8 : pos + 12])[0]
        record_total = 16 + incl_len
        if pos + record_total > len(data):
            sys.exit(f"file ends mid-record-payload at offset {pos} (declared len {incl_len}): {path}")
        pos += record_total
        offsets.append(pos)
    return offsets


def main():
    if len(sys.argv) != 3:
        sys.exit(f"usage: {sys.argv[0]} <file> <n_segments>")
    path = sys.argv[1]
    n = int(sys.argv[2])
    if n < 1:
        sys.exit("n_segments must be >= 1")

    offsets = record_offsets(path)  # offsets[0]=24 (after header), offsets[-1]=file size
    n_records = len(offsets) - 1
    if n_records < n:
        sys.exit(f"file only has {n_records} records, can't split into {n} segments")

    # n+1 cumulative-size boundaries, evenly spaced by record index, from
    # "header only" (offsets[0] == 24) up to the whole file (offsets[-1]).
    for i in range(0, n + 1):
        record_idx = (n_records * i) // n
        print(offsets[record_idx])


if __name__ == "__main__":
    main()