#!/usr/bin/env python3
"""
Pulls a capture off the ITS-G5 receiver's in-memory pcap buffer over the existing USB serial
connection and saves it as a real .pcap file on this machine.

Requires the firmware to be built with:
    Example Configuration -> Select destination to store pcap file -> Memory

Usage (typical):
    1. Close idf.py monitor (only one program can hold the COM port at a time).
    2. Run a capture on the device: `sniffer -P` ... let it run ... `sniffer --stop`
    3. python dump_pcap.py COM5

The device has no access to this computer's filesystem, so it can't write here directly. Instead,
`pcap --dump` streams the raw pcap bytes back over the same serial link, wrapped in plain-text
markers ("===PCAP-DUMP-START:<len>===" ... raw bytes ... "===PCAP-DUMP-END==="). This script finds
those markers and writes just the raw bytes out as a .pcap file.

Install dependency once: pip install pyserial
"""

import argparse
import datetime
import re
import sys

try:
    import serial
except ImportError:
    print("Missing dependency. Install it with:  pip install pyserial", file=sys.stderr)
    sys.exit(1)

START_RE = re.compile(rb"===PCAP-DUMP-START:(\d+)===\n")
END_MARKER = b"\n===PCAP-DUMP-END===\n"


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("port", help="Serial port the device is on, e.g. COM5")
    parser.add_argument("-b", "--baud", type=int, default=115200, help="Baud rate (default: 115200)")
    parser.add_argument("-o", "--outdir", default="recordings", help="Output directory (default: ./recordings)")
    parser.add_argument("-t", "--timeout", type=float, default=15.0, help="Seconds to wait for the dump to start")
    args = parser.parse_args()

    import os
    os.makedirs(args.outdir, exist_ok=True)

    print(f"Opening {args.port} @ {args.baud}...")
    with serial.Serial(args.port, args.baud, timeout=1) as ser:
        # Nudge the console in case there's stale input, then request the dump.
        ser.reset_input_buffer()
        ser.write(b"\r\n")
        ser.write(b"pcap -f dump --dump\r\n")

        print("Waiting for dump to start...")
        buf = b""
        match = None
        deadline = datetime.datetime.now() + datetime.timedelta(seconds=args.timeout)
        while datetime.datetime.now() < deadline:
            chunk = ser.read(256)
            if chunk:
                buf += chunk
                match = START_RE.search(buf)
                if match:
                    break
        if not match:
            print("Timed out waiting for '===PCAP-DUMP-START:...===' marker.\n"
                  "Check that: the firmware is built with the Memory pcap destination, a capture was\n"
                  "actually taken ('sniffer -P' then 'sniffer --stop'), and no other program (like\n"
                  "idf.py monitor) is holding the serial port open.", file=sys.stderr)
            sys.exit(1)

        length = int(match.group(1))
        print(f"Dump starting, {length} bytes expected.")

        # Anything after the marker in our buffer is already part of the payload.
        payload = buf[match.end():]
        remaining = length - len(payload)
        while remaining > 0:
            chunk = ser.read(min(remaining, 4096))
            if not chunk:
                print(f"Serial read timed out with {remaining} bytes still missing.", file=sys.stderr)
                sys.exit(1)
            payload += chunk
            remaining -= len(chunk)

        # Drain (and sanity-check) the trailing end marker, but don't fail hard if it's not exact.
        tail = ser.read(len(END_MARKER))
        if tail != END_MARKER:
            print("Warning: end marker didn't match exactly - payload may still be fine.", file=sys.stderr)

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    outpath = os.path.join(args.outdir, f"capture_{timestamp}.pcap")
    with open(outpath, "wb") as f:
        f.write(payload)

    print(f"Saved {len(payload)} bytes to {outpath}")


if __name__ == "__main__":
    main()