#!/usr/bin/env python3
"""
Continuously listens on the ITS-G5 receiver's serial console and writes every captured packet into a
live-growing .pcap file, with no console commands needed on the device side.

The firmware streams every packet it captures out over the same serial connection the console runs on,
automatically, as soon as the sniffer is running (which happens on boot by default). Each packet is framed
with plain-text markers so this script can pull the binary pcap bytes out of the stream even though
regular log lines are interleaved with it:

    ===PCAP-LIVE-HEADER:<len>===\\n<24 raw bytes>\\n   (sent once, the pcap global header)
    ===PCAP-LIVE-PKT:<len>===\\n<raw bytes>\\n          (sent once per captured packet)

Usage:
    python live_capture.py COM5

Runs until you press Ctrl+C. Writes to recordings/capture_<timestamp>.pcap, flushing after every packet
so you can open the file in Wireshark while it's still being written (use "File > Open" again, or
Wireshark's own "Follow" won't auto-refresh but re-opening will show the latest packets).

Install dependency once: pip install pyserial
"""

import argparse
import datetime
import os
import re
import sys
import time

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

# \r? because the ESP console emits CRLF: on Windows the markers arrive as
# "===PCAP-LIVE-PKT:310===\r\n", which never matched a bare \n and left the capture silently
# empty while the device was streaming perfectly well.
HEADER_RE = re.compile(rb"===PCAP-LIVE-HEADER:(\d+)===\r?\n")
PKT_RE = re.compile(rb"===PCAP-LIVE-PKT:(\d+)===\r?\n")

LINKTYPE_ETHERNET = 1
LINKTYPE_IEEE802_11_RADIOTAP = 127


def mac_str(b):
    return ":".join(f"{x:02x}" for x in b)


def build_default_pcap_header(link_type):
    """Synthesizes the same 24-byte global pcap header the firmware would have sent, for when we
    connect after the device's one-time header already went out (see the race note in main())."""
    header = bytearray(24)
    header[0:4] = bytes([0xD4, 0xC3, 0xB2, 0xA1])  # magic (LE bytes of 0xA1B2C3D4)
    header[4:6] = (2).to_bytes(2, "little")         # major version
    header[6:8] = (4).to_bytes(2, "little")         # minor version
    header[16:20] = (0x40000).to_bytes(4, "little")  # snaplen
    header[20:24] = link_type.to_bytes(4, "little")
    return bytes(header)


def summarize_packet(link_type, record_bytes, index):
    """Best-effort human-readable one-line summary of a captured packet, for live feedback.
    record_bytes is the raw 16-byte pcap record header followed by the captured frame."""
    seconds = int.from_bytes(record_bytes[0:4], "little")
    microseconds = int.from_bytes(record_bytes[4:8], "little")
    cap_len = int.from_bytes(record_bytes[8:12], "little")
    frame = record_bytes[16:]
    ts = f"{seconds}.{microseconds:06d}"

    if link_type == LINKTYPE_IEEE802_11_RADIOTAP and len(frame) >= 24:
        radiotap_len = int.from_bytes(frame[2:4], "little")
        rssi = frame[8] - 256 if frame[8] >= 128 else frame[8]
        station_id = int.from_bytes(frame[16:24], "little")
        mac_frame = frame[radiotap_len:]
        if len(mac_frame) >= 16:
            dst = mac_str(mac_frame[4:10])
            src = mac_str(mac_frame[10:16])
        else:
            dst = src = "?"
        station = f"{station_id:012x}" if station_id else "unknown"
        return (f"#{index:<5} [{ts}] len={cap_len:<5} rssi={rssi:>4}dBm  "
                f"station={station}  {src} -> {dst}")

    if link_type == LINKTYPE_ETHERNET and len(frame) >= 14:
        dst = mac_str(frame[0:6])
        src = mac_str(frame[6:12])
        ethertype = int.from_bytes(frame[12:14], "big")
        return f"#{index:<5} [{ts}] len={cap_len:<5} eth  {src} -> {dst}  type=0x{ethertype:04x}"

    return f"#{index:<5} [{ts}] len={cap_len:<5} (unrecognized frame format)"


def undo_crlf(chunk, state):
    """Undo the CR the device console inserts before every LF.

    ESP-IDF's newlib console converts LF to CRLF on its way out, and that happens to every 0x0A
    byte of the binary pcap stream too, not only to log text. Each inserted CR shifts everything
    after it, so pcap record headers and captured frames alike come out corrupt. This is the
    "byte inserted mid-frame" seen in older recordings; with DENM traffic on air it wrecks most
    of a capture (measured 2026-09-14: a 787 KB file parsed cleanly for only 82 records).

    Dropping one CR immediately before each LF undoes it exactly, provided it is done on the raw
    stream before any framing and a trailing CR is carried across read boundaries. CR and LF are
    written as byte values here so the transformation cannot be confused with an escape.
    """
    CR, LF = bytes([13]), bytes([10])
    if state["pending_cr"]:
        chunk = CR + chunk
        state["pending_cr"] = False
    if chunk.endswith(CR):
        chunk = chunk[:-1]
        state["pending_cr"] = True
    return chunk.replace(CR + LF, LF)


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)")
    args = parser.parse_args()

    os.makedirs(args.outdir, exist_ok=True)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    outpath = os.path.join(args.outdir, f"capture_{timestamp}.pcap")

    print(f"Opening {args.port} @ {args.baud}...")
    print(f"Writing live capture to {outpath}")
    print("Press Ctrl+C to stop.")

    header_written = False
    link_type = None
    packet_count = 0
    buf = b""
    total_bytes = 0
    last_status = time.monotonic()
    printed_raw_preview = False

    crlf_state = {"pending_cr": False}

    with serial.Serial(args.port, args.baud, timeout=1) as ser, open(outpath, "wb") as outfile:
        def read_bytes(n):
            return undo_crlf(ser.read(n), crlf_state)

        try:
            while True:
                chunk = read_bytes(256)
                if chunk:
                    buf += chunk
                    total_bytes += len(chunk)

                now = time.monotonic()
                if now - last_status >= 2:
                    last_status = now
                    print(f"[diagnostic] {total_bytes} raw bytes received so far, "
                          f"{packet_count} packets recognized, header_written={header_written}")
                    if total_bytes > 0 and not printed_raw_preview and not header_written and not PKT_RE.search(buf):
                        # We're getting bytes but none of them look like our markers - show a preview
                        # so we can tell whether this is plain log text (markers just haven't shown up
                        # yet), garbage (baud/port mismatch), or something else entirely.
                        preview = buf[:200]
                        print(f"[diagnostic] no markers matched yet - raw preview: {preview!r}")
                        printed_raw_preview = True
                    elif total_bytes == 0:
                        print("[diagnostic] zero bytes received from the port at all - this points at "
                              "the wrong COM port, another program holding the port, or a port that "
                              "isn't actually wired to the console/sniffer output.")

                # The device only sends the global header once, right when the sniffer first starts
                # (typically within a second or two of boot). If this script connects even slightly
                # late - very likely right after a fresh flash, since esptool itself resets the board -
                # that header is already gone before we ever see it. Rather than blocking forever
                # waiting for a header that's never coming, look for whichever marker shows up first.
                header_match = None if header_written else HEADER_RE.search(buf)
                pkt_match = PKT_RE.search(buf)

                if header_match and (not pkt_match or header_match.start() < pkt_match.start()):
                    length = int(header_match.group(1))
                    buf = buf[header_match.end():]
                    while len(buf) < length:
                        buf += read_bytes(length - len(buf))
                    header_bytes = buf[:length]
                    outfile.write(header_bytes)
                    outfile.flush()
                    buf = buf[length:]
                    header_written = True
                    if length >= 24:
                        link_type = int.from_bytes(header_bytes[20:24], "little")
                    print(f"Got pcap global header (link type {link_type}) - device is streaming.\n")
                    continue

                if not header_written and pkt_match:
                    link_type = LINKTYPE_IEEE802_11_RADIOTAP
                    outfile.write(build_default_pcap_header(link_type))
                    outfile.flush()
                    header_written = True
                    print("Note: missed the device's one-time pcap header (it was likely sent before "
                          "this script connected, e.g. right after a flash/reset) - assuming WLAN "
                          "radiotap capture and writing a default header instead.\n")
                    # fall through and process pkt_match below, don't discard this packet

                if not pkt_match:
                    # Keep the buffer from growing unbounded while waiting for a marker, but don't
                    # discard anything - a marker could be split across reads.
                    if len(buf) > 65536:
                        buf = buf[-4096:]
                    continue

                length = int(pkt_match.group(1))
                buf = buf[pkt_match.end():]
                while len(buf) < length:
                    buf += read_bytes(length - len(buf))
                record_bytes = buf[:length]
                outfile.write(record_bytes)
                outfile.flush()
                buf = buf[length:]
                packet_count += 1
                print(summarize_packet(link_type, record_bytes, packet_count))
        except KeyboardInterrupt:
            print(f"\nStopped. {packet_count} packets saved to {outpath}")


if __name__ == "__main__":
    main()