#!/usr/bin/env python3
"""vtcmd.py -- talk to the LANCE command payload (harness/udpstream/lance_cmd.c).

The payload runs entirely on its own (no OS stack), so this is the only way to
inspect the machine while it owns the chip.

Protocol, ethertype 0x88B6, both directions:

    [14] 'K'  [15] 'S'  [16] cmd
    cmd 1  read : [17..20] address (32-bit big-endian BIT address), [21..22] len
    cmd 2  stats: no arguments
    cmd 3  poke : [17..20] address (BIT), [21..22] len, [23..] the bytes
    cmd 4  jump : [17..20] address (BIT) -- reply is sent, then control moves
    cmd 5  reboot: no arguments -- jumps to the ROM reset vector
    reply      : [16] cmd echoed, [17..18] status, [19..] data

Every address in this protocol is a BIT address, the way the TMS34010 sees
memory.  For a byte address, send byte_addr * 8.

Usage:

    # one-off statistics (frames/bytes the payload has seen)
    sudo python3 vtcmd.py stats --iface vt1200

    # write 4 bytes at bit address 0x018D0300 (byte 0x31A60)
    sudo python3 vtcmd.py poke --iface vt1200 --addr 0x018D0300 --data "de ad be ef"

    # run code at an address, or soft-reset the machine back into the firmware
    sudo python3 vtcmd.py jump   --iface vt1200 --addr 0x018C0000
    sudo python3 vtcmd.py reboot --iface vt1200

    # read 32 bytes of memory at bit address 0x018D0000
    sudo python3 vtcmd.py read --iface vt1200 --addr 0x018D0000 --len 32

    # throughput: blast N frames, then ask the payload how many it counted
    sudo python3 vtcmd.py bench --iface vt1200 --count 4000 --size 1514
"""
import argparse
import socket
import struct
import sys
import time

ETHCMD = 0x88B6
VT_MAC = bytes.fromhex("08002b164b5a")


def parse_mac(s):
    return bytes(int(p, 16) for p in s.split(":"))


def iface_mac(iface):
    with open("/sys/class/net/%s/address" % iface) as f:
        return parse_mac(f.read().strip())


def open_sock(iface):
    s = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(ETHCMD))
    s.bind((iface, 0))
    return s


def frame(iface, cmd, addr=0, length=0, size=60, payload=b""):
    src = iface_mac(iface)
    body = bytes([ord('K'), ord('S'), cmd]) + struct.pack(">I", addr) + \
        struct.pack(">H", length) + payload
    body = body + bytes(max(0, size - 14 - len(body)))
    return b"\xff" * 6 + src + struct.pack(">H", ETHCMD) + body


def request(sock, iface, cmd, addr=0, length=0, timeout=2.0, tries=3,
            payload=b""):
    """Send a command and wait for its reply.

    Retried because the first frame after the payload takes the chip over is
    routinely lost: the ring is armed slightly before the receive loop runs.
    """
    for attempt in range(tries):
        sock.send(frame(iface, cmd, addr, length,
                        size=max(60, 23 + len(payload)), payload=payload))
        end = time.time() + timeout
        sock.settimeout(0.2)
        while time.time() < end:
            try:
                data, addr_in = sock.recvfrom(2048)
            except socket.timeout:
                continue
            if len(data) < 19:
                continue
            if data[12:14] != struct.pack(">H", ETHCMD):
                continue
            if data[14:16] != b"KS":
                continue
            if data[16] != cmd:
                continue
            status = (data[17] << 8) | data[18]
            return status, data[19:]
    return None, None


def cmd_stats(sock, iface, quiet=False):
    status, data = request(sock, iface, 2)
    if status is None:
        if not quiet:
            print("no reply -- is the payload running and the link up?")
        return None
    rxf, nbytes, cmds, chaindrop, txfull, errbits = struct.unpack(">IIHHHH", data[:16])
    flags = []
    if errbits & 0x1000:
        flags.append("MISS(no descriptor)")
    if errbits & 0x0800:
        flags.append("MERR")
    if errbits & 0x4000:
        flags.append("BABL")
    if errbits & 0x2000:
        flags.append("CERR")
    if errbits & 0x8000:
        flags.append("ERR")
    if not quiet:
        print("payload: RXF=%d  bytes=%d  cmds=%d  chained_dropped=%d  txfull=%d  "
              "csr0_err=%#06x%s  status=%#06x"
              % (rxf, nbytes, cmds, chaindrop, txfull, errbits,
                 (" [" + ",".join(flags) + "]") if flags else " [none]", status))
    return rxf, nbytes, cmds, chaindrop, txfull, errbits


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("cmd", choices=["stats", "read", "poke", "jump",
                                    "reboot", "bench"])
    ap.add_argument("--iface", default="vt1200")
    ap.add_argument("--addr", type=lambda v: int(v, 0), default=0)
    ap.add_argument("--len", type=int, default=32)
    ap.add_argument("--data", default="", help="poke: bytes, e.g. 'de ad be ef'")
    ap.add_argument("--file", default="", help="poke: read the bytes from a file")
    ap.add_argument("--count", type=int, default=0,
                    help="send exactly N frames instead of a timed burst")
    ap.add_argument("--seconds", type=float, default=3.0,
                    help="duration of the timed burst when --count is 0")
    ap.add_argument("--size", type=int, default=1514)
    ap.add_argument("--etype", type=lambda v: int(v, 0), default=0x88B7,
                    help="ethertype for benchmark frames (must NOT be 0x88B6, "
                         "or the payload answers every one of them)")
    ap.add_argument("--dst", default="ff:ff:ff:ff:ff:ff",
                    help="destination MAC; unicast avoids the switch's "
                         "broadcast storm control")
    ap.add_argument("--rate", type=float, default=0,
                    help="pace to at most N frames/s (0 = as fast as possible)")
    ap.add_argument("--busy", action="store_true",
                    help="busy-wait instead of time.sleep while pacing; "
                         "sleep granularity is ~250us and bunches frames "
                         "into microbursts that the switch drops")
    ap.add_argument("--terse", action="store_true",
                    help="one CSV line per run, for sweeps and for the report")
    ap.add_argument("--timeout", type=float, default=2.0)
    a = ap.parse_args()

    sock = open_sock(a.iface)

    if a.cmd == "stats":
        return 0 if cmd_stats(sock, a.iface) else 1

    if a.cmd == "read":
        status, data = request(sock, a.iface, 1, a.addr, a.len, a.timeout)
        if status is None:
            print("no reply")
            return 1
        print("read %#010x len %d status %#06x:" % (a.addr, a.len, status))
        for i in range(0, min(len(data), a.len), 16):
            print("  %04x  %s" % (i, data[i:i + 16].hex(" ")))
        return 0

    if a.cmd == "poke":
        if a.file:
            with open(a.file, "rb") as f:
                blob = f.read()
        else:
            blob = bytes(int(b, 16) for b in a.data.replace(",", " ").split())
        if not blob:
            print("nothing to poke: use --data or --file")
            return 1
        status, data = request(sock, a.iface, 3, a.addr, len(blob), a.timeout,
                               payload=blob)
        if status is None:
            print("no reply")
            return 1
        print("poked %d bytes at bit %#010x (byte %#x) status %#06x"
              % (len(blob), a.addr, a.addr >> 3, status))
        return 0

    if a.cmd in ("jump", "reboot"):
        cmd = 4 if a.cmd == "jump" else 5
        status, data = request(sock, a.iface, cmd, a.addr, 0, a.timeout)
        if status is None:
            print("no reply")
            return 1
        target = struct.unpack(">I", data[:4])[0] if len(data) >= 4 else 0
        print("%s: control -> bit %#010x (byte %#x), status %#06x"
              % (a.cmd, target, target >> 3, status))
        print("the payload has handed control over; it will not answer again")
        return 0

    # bench: offer frames continuously for --seconds, then let the payload's
    # counter settle.  A fixed-duration burst makes the window long enough that
    # the settle time is negligible, and retrying on ENOBUFS paces the host to
    # whatever its socket buffer allows instead of aborting.
    before = cmd_stats(sock, a.iface, quiet=a.terse)
    size = max(a.size, 60)
    body = bytes(size - 14)
    f = parse_mac(a.dst) + iface_mac(a.iface) + struct.pack(">H", a.etype) + body
    sent = 0
    t0 = time.time()
    if a.count:
        for i in range(a.count):
            try:
                sock.send(f)
                sent += 1
            except OSError:
                time.sleep(0.001)
                sock.send(f)
                sent += 1
    else:
        while time.time() - t0 < a.seconds:
            if a.rate:
                want = sent / a.rate
                behind = want - (time.time() - t0)
                if behind > 0:
                    if a.busy:
                        end = time.time() + behind
                        while time.time() < end:
                            pass
                    else:
                        time.sleep(behind)
            try:
                sock.send(f)
                sent += 1
            except OSError:
                time.sleep(0.001)
    t1 = time.time()

    # poll until the *data* counter is stable (each stats request is itself a
    # received command frame, so comparing rxf alone would never settle)
    after = before
    quiet = 0
    last = before[0] - before[2]
    deadline = time.time() + 15.0
    while time.time() < deadline:
        time.sleep(0.2)
        now = cmd_stats(sock, a.iface, quiet=a.terse)
        if not now:
            break
        after = now
        cur = now[0] - now[2]
        if cur == last:
            quiet += 1
            if quiet >= 2:
                break
        else:
            quiet = 0
        last = cur
    t2 = time.time()

    if before and after and after[0] >= before[0]:
        dt = t1 - t0
        dcmds = after[2] - before[2]          # our own stats frames
        frames = (after[0] - before[0]) - dcmds
        nbytes = (after[1] - before[1]) - 60 * dcmds
        if a.terse:
            print("%d,%g,%.2f,%d,%d,%d,%.0f,%.1f,%.2f,%.1f,%#06x,%d,%d"
                  % (size, a.rate, dt, sent, frames, nbytes,
                     frames / dt if dt else 0,
                     nbytes / 1024.0 / dt if dt else 0,
                     nbytes * 8 / dt / 1e6 if dt else 0,
                     100.0 * frames / sent if sent else 0,
                     after[5], after[4], after[3]))
            return 0
        print("sent %d frames of %d B over %.3f s (%.0f frames/s offered)"
              % (sent, size, dt, sent / dt if dt else 0))
        print("payload received %d frames (%d bytes) while they were offered"
              % (frames, nbytes))
        print("  receive rate: %.0f frames/s, %.1f KiB/s, %.1f Mbit/s"
              % (frames / dt if dt else 0,
                 nbytes / 1024.0 / dt if dt else 0,
                 nbytes * 8 / dt / 1e6 if dt else 0))
        print("  delivery ratio: %.1f%%  (below 100%% = the ring or the wire "
              "could not keep up)" % (100.0 * frames / sent if sent else 0))
        print("  drain after the burst: %.2f s" % max(0.0, t2 - t1))
    return 0


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