#!/usr/bin/env python3
"""Extract the host -> VT1200 data stream from a LAT capture.

Reads a pcap, keeps frames *from* the host *to* the terminal, decodes the LAT
header/slots (same layout as lat_probe.py), and concatenates the Data_a slot
payloads in order.  The result is the exact byte stream the terminal received.

    python3 extract_stream.py filter-test2.pcap -o stream.bin
    python3 extract_stream.py filter-test2.pcap --stats
"""
import argparse
import struct
import sys

HOST_MAC = "2c:16:db:a0:e7:b1"
VT1200_MAC = "08:00:2b:16:4b:5a"
ETHERTYPE_LAT = 0x6004
SLOT_DATA_A = 0x0


def mac_bytes(s):
    return bytes(int(x, 16) for x in s.split(":"))


def read_pcap(path):
    """Yield (timestamp, frame_bytes) from a libpcap file (both endians)."""
    with open(path, "rb") as fh:
        hdr = fh.read(24)
        if len(hdr) < 24:
            return
        magic = hdr[:4]
        if magic in (b"\xd4\xc3\xb2\xa1", b"\x4d\x3c\xb2\xa1"):
            endian = "<"
        elif magic in (b"\xa1\xb2\xc3\xd4", b"\xa1\xb2\x3c\x4d"):
            endian = ">"
        else:
            raise ValueError(f"{path}: not a libpcap file (magic {magic.hex()})")
        while True:
            ph = fh.read(16)
            if len(ph) < 16:
                return
            ts_sec, ts_usec, caplen, _orig = struct.unpack(endian + "IIII", ph)
            data = fh.read(caplen)
            if len(data) < caplen:
                return
            yield ts_sec + ts_usec / 1e6, data


def parse_lat(frame):
    """Return (src_mac, slots) for a LAT frame, else None."""
    if len(frame) < 14 or struct.unpack_from(">H", frame, 12)[0] != ETHERTYPE_LAT:
        return None
    src = frame[6:12]
    lat = frame[14:]
    if len(lat) < 8:
        return None
    nbr = lat[1]
    slots = []
    off = 8
    for _ in range(nbr):
        if off + 4 > len(lat):
            break
        ds, ss, bc, ct = lat[off], lat[off + 1], lat[off + 2], lat[off + 3]
        slots.append((ds, ss, bc, ct >> 4, ct & 0xF, bytes(lat[off + 4:off + 4 + bc])))
        off += 4 + bc
    return src, slots


def extract(pcap_path, host=HOST_MAC, terminal=VT1200_MAC):
    """Return (stream, events) where events = [(timestamp, length), ...]."""
    stream = bytearray()
    events = []
    want = mac_bytes(host)
    vmac = mac_bytes(terminal)
    for ts, frame in read_pcap(pcap_path):
        got = parse_lat(frame)
        if not got:
            continue
        src, slots = got
        if src != want:
            continue                       # only host -> terminal
        for (_ds, _ss, _bc, _cr, st, data) in slots:
            if st == SLOT_DATA_A and data:
                stream += data
                events.append((ts, len(data)))
    return bytes(stream), events


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pcap")
    ap.add_argument("-o", "--out")
    ap.add_argument("--stats", action="store_true")
    ap.add_argument("--host", default=HOST_MAC)
    ap.add_argument("--terminal", default=VT1200_MAC)
    args = ap.parse_args(argv)

    stream, events = extract(args.pcap, args.host, args.terminal)
    print(f"{args.pcap}: {len(stream)} bytes host->terminal in {len(events)} slots")
    if events:
        span = events[-1][0] - events[0][0]
        print(f"  first {events[0][0]:.3f}s  last {events[-1][0]:.3f}s  span {span:.3f}s")
    print(f"  ESC present: {0x1b in stream}   high-bit bytes: "
          f"{sum(1 for b in stream if b > 0x7f)}")
    if args.stats:
        for ts, n in events:
            print(f"  {ts:.3f}  +{n}")
    if args.out:
        with open(args.out, "wb") as fh:
            fh.write(stream)
        print(f"wrote {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())