Vibed advanced python test script
This commit is contained in:
Your Name
@@ -0,0 +1,346 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Send binary commands to a device over a serial port, optionally receiving and/or requiring an ACK."""
|
||||
|
||||
import argparse
|
||||
import struct
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
||||
import serial
|
||||
|
||||
DEFAULT_PORT = "/dev/ttyACM0"
|
||||
DEFAULT_BAUD = 115200
|
||||
DEFAULT_DEVICE_ID = 0
|
||||
|
||||
COMMAND_PREFIX = 0x69
|
||||
|
||||
COMMAND_ACK = 0
|
||||
COMMAND_NACK = 1
|
||||
|
||||
# name -> command id
|
||||
COMMANDS = {
|
||||
"ACK": 0,
|
||||
"NACK": 1,
|
||||
"LED_TOGGLE": 2,
|
||||
"LED_SET": 3,
|
||||
"SERVO_SET": 4,
|
||||
"SERVO_SET_ALL": 5,
|
||||
"ADC_READ_RAW": 6,
|
||||
"ADC_READ_RAW_ALL": 7,
|
||||
"ADC_READ": 8,
|
||||
"ADC_READ_ALL": 9,
|
||||
"ADC_SET_READ": 10,
|
||||
}
|
||||
COMMAND_NAMES = {v: k for k, v in COMMANDS.items()}
|
||||
|
||||
# Commands for which the data argument is OPTIONAL.
|
||||
# Every command not listed here REQUIRES data.
|
||||
# >>> Edit this list to match your firmware. <<<
|
||||
COMMANDS_DATA_OPTIONAL = {
|
||||
"ACK",
|
||||
"NACK",
|
||||
"LED_TOGGLE",
|
||||
"ADC_READ_RAW_ALL",
|
||||
"ADC_READ_ALL",
|
||||
}
|
||||
|
||||
# struct formats for typed data tokens (all little-endian)
|
||||
TYPE_FORMATS = {
|
||||
"u8": "<B", "i8": "<b",
|
||||
"u16": "<H", "i16": "<h",
|
||||
"u32": "<I", "i32": "<i",
|
||||
"f32": "<f",
|
||||
}
|
||||
|
||||
# Exit codes
|
||||
EXIT_OK = 0
|
||||
EXIT_ERROR = 1
|
||||
EXIT_NO_ACK = 2
|
||||
EXIT_NACK = 3
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Packet helpers
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
def calculate_crc(msg: bytes) -> int:
|
||||
s = sum(msg) & 0xFF
|
||||
return (-s) & 0xFF
|
||||
|
||||
|
||||
def make_packet(command: int, data: bytes = b"", device_id: int = DEFAULT_DEVICE_ID) -> bytes:
|
||||
pkt = bytearray()
|
||||
pkt.append(COMMAND_PREFIX)
|
||||
pkt.append(len(data))
|
||||
pkt.append(device_id)
|
||||
pkt.append(command)
|
||||
pkt.append(0) # CRC placeholder
|
||||
pkt.extend(data)
|
||||
|
||||
pkt[4] = calculate_crc(pkt[:4] + pkt[5:])
|
||||
return bytes(pkt)
|
||||
|
||||
|
||||
def verify_crc(packet: bytes) -> bool:
|
||||
return packet[4] == calculate_crc(packet[:4] + packet[5:])
|
||||
|
||||
|
||||
def packet_size(buf: bytes):
|
||||
if len(buf) < 2:
|
||||
return None
|
||||
return 5 + buf[1]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Receiver thread
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
class Receiver(threading.Thread):
|
||||
"""Parses incoming packets / ASCII logs. Prints them only if verbose."""
|
||||
|
||||
def __init__(self, ser, verbose):
|
||||
super().__init__(daemon=True)
|
||||
self.ser = ser
|
||||
self.verbose = verbose
|
||||
self.stop_event = threading.Event()
|
||||
self.ack_event = threading.Event()
|
||||
self.result = None # "ACK" or "NACK" once received
|
||||
|
||||
def _print(self, *args):
|
||||
if self.verbose:
|
||||
print(*args)
|
||||
|
||||
def run(self):
|
||||
rx = bytearray()
|
||||
|
||||
while not self.stop_event.is_set():
|
||||
data = self.ser.read(64)
|
||||
if not data:
|
||||
continue
|
||||
|
||||
rx.extend(data)
|
||||
|
||||
while rx:
|
||||
if rx[0] == COMMAND_PREFIX:
|
||||
size = packet_size(rx)
|
||||
if size is None or len(rx) < size:
|
||||
break
|
||||
|
||||
pkt = bytes(rx[:size])
|
||||
del rx[:size]
|
||||
|
||||
if not verify_crc(pkt):
|
||||
self._print("RX: Bad CRC:", pkt.hex(" "))
|
||||
continue
|
||||
|
||||
length = pkt[1]
|
||||
dev_id = pkt[2]
|
||||
cmd = pkt[3]
|
||||
|
||||
if cmd == COMMAND_ACK:
|
||||
self._print(f"<-- ACK (device={dev_id})")
|
||||
self.result = "ACK"
|
||||
self.ack_event.set()
|
||||
elif cmd == COMMAND_NACK:
|
||||
self._print(f"<-- NACK (device={dev_id})")
|
||||
self.result = "NACK"
|
||||
self.ack_event.set()
|
||||
else:
|
||||
name = COMMAND_NAMES.get(cmd, str(cmd))
|
||||
self._print(f"<-- Command {name} len={length}")
|
||||
self._print(f"<-- RX: {pkt.hex(' ')}")
|
||||
else:
|
||||
idx = rx.find(b"\n")
|
||||
if idx == -1:
|
||||
break
|
||||
|
||||
line = rx[:idx + 1]
|
||||
del rx[:idx + 1]
|
||||
|
||||
try:
|
||||
self._print("[LOG]", line.decode().rstrip())
|
||||
except UnicodeDecodeError:
|
||||
self._print("[RAW]", line.hex())
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Argument parsing
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
def parse_command(text: str):
|
||||
"""Accept a command name (case-insensitive) or a number. Returns (id, name)."""
|
||||
key = text.upper()
|
||||
if key in COMMANDS:
|
||||
return COMMANDS[key], key
|
||||
|
||||
try:
|
||||
val = int(text, 0)
|
||||
except ValueError:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"unknown command '{text}' (use --list-commands to see them)"
|
||||
)
|
||||
if not 0 <= val <= 255:
|
||||
raise argparse.ArgumentTypeError("command number must be 0-255")
|
||||
return val, COMMAND_NAMES.get(val, f"#{val}")
|
||||
|
||||
|
||||
def parse_data_token(tok: str) -> bytes:
|
||||
"""
|
||||
Formats:
|
||||
5 -> u8
|
||||
u16:500 -> unsigned 16-bit little-endian
|
||||
i32:-3 -> signed 32-bit little-endian
|
||||
f32:1.5 -> float32 little-endian
|
||||
hex:0a0b0c -> raw bytes
|
||||
"""
|
||||
if ":" in tok:
|
||||
kind, value = tok.split(":", 1)
|
||||
kind = kind.lower()
|
||||
else:
|
||||
kind, value = "u8", tok
|
||||
|
||||
if kind == "hex":
|
||||
return bytes.fromhex(value)
|
||||
|
||||
fmt = TYPE_FORMATS.get(kind)
|
||||
if fmt is None:
|
||||
raise ValueError(f"unknown type '{kind}' (use u8/i8/u16/i16/u32/i32/f32/hex)")
|
||||
|
||||
num = float(value) if kind == "f32" else int(value, 0)
|
||||
return struct.pack(fmt, num)
|
||||
|
||||
|
||||
def build_parser():
|
||||
commands_help = ", ".join(f"{name}={cid}" for name, cid in COMMANDS.items())
|
||||
optional_help = ", ".join(sorted(COMMANDS_DATA_OPTIONAL)) or "(none)"
|
||||
|
||||
p = argparse.ArgumentParser(
|
||||
description="Send a binary command to a serial device.",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=f"""\
|
||||
commands:
|
||||
{commands_help}
|
||||
|
||||
commands where data is optional:
|
||||
{optional_help}
|
||||
|
||||
data tokens (all multi-byte values are little-endian):
|
||||
5 u8 (default type)
|
||||
u16:500 also: u8 i8 u16 i16 u32 i32 f32
|
||||
hex:0a0b raw bytes
|
||||
|
||||
examples:
|
||||
%(prog)s LED_TOGGLE
|
||||
%(prog)s ADC_SET_READ u32:1 --require-ack
|
||||
%(prog)s SERVO_SET 0 u16:1500 --rx --rx-time 5
|
||||
%(prog)s -p /dev/ttyUSB0 -b 9600 ADC_READ_ALL --rx
|
||||
""",
|
||||
)
|
||||
|
||||
p.add_argument("command", nargs="?", type=parse_command,
|
||||
help="command name or number")
|
||||
p.add_argument("data", nargs="*", help="command data tokens (see below)")
|
||||
|
||||
p.add_argument("-p", "--port", default=DEFAULT_PORT,
|
||||
help="serial port (default: %(default)s)")
|
||||
p.add_argument("-b", "--baud", type=int, default=DEFAULT_BAUD,
|
||||
help="baud rate (default: %(default)s)")
|
||||
p.add_argument("-d", "--device-id", type=lambda s: int(s, 0), default=DEFAULT_DEVICE_ID,
|
||||
help="device id in the packet header (default: %(default)s)")
|
||||
|
||||
p.add_argument("--rx", action="store_true",
|
||||
help="receive and print incoming packets/logs (default: off)")
|
||||
p.add_argument("--rx-time", type=float, default=10.0,
|
||||
help="seconds to keep listening after sending, with --rx (default: %(default)s)")
|
||||
|
||||
p.add_argument("--require-ack", action="store_true",
|
||||
help="wait for an ACK; exit 2 on timeout, 3 on NACK (default: off)")
|
||||
p.add_argument("--ack-timeout", type=float, default=1.0,
|
||||
help="seconds to wait for ACK (default: %(default)s)")
|
||||
|
||||
p.add_argument("--list-commands", action="store_true",
|
||||
help="list known commands and exit")
|
||||
return p
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Main
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.list_commands:
|
||||
for name, cid in COMMANDS.items():
|
||||
opt = "data optional" if name in COMMANDS_DATA_OPTIONAL else "data required"
|
||||
print(f"{cid:3d} {name:<18} {opt}")
|
||||
return EXIT_OK
|
||||
|
||||
if args.command is None:
|
||||
parser.error("the following arguments are required: command")
|
||||
|
||||
cmd_id, cmd_name = args.command
|
||||
|
||||
if not args.data and cmd_name not in COMMANDS_DATA_OPTIONAL:
|
||||
parser.error(f"command {cmd_name} requires data")
|
||||
|
||||
data = bytearray()
|
||||
try:
|
||||
for tok in args.data:
|
||||
data.extend(parse_data_token(tok))
|
||||
except (ValueError, struct.error) as e:
|
||||
parser.error(f"bad data: {e}")
|
||||
|
||||
if len(data) > 255:
|
||||
parser.error(f"data too long ({len(data)} bytes, max 255)")
|
||||
|
||||
packet = make_packet(cmd_id, bytes(data), args.device_id)
|
||||
|
||||
try:
|
||||
ser = serial.Serial(args.port, args.baud, timeout=0.05)
|
||||
except serial.SerialException as e:
|
||||
print(f"ERROR: could not open {args.port}: {e}", file=sys.stderr)
|
||||
return EXIT_ERROR
|
||||
|
||||
receiver = None
|
||||
exit_code = EXIT_OK
|
||||
|
||||
try:
|
||||
# Start the receiver BEFORE sending so a fast ACK isn't missed.
|
||||
if args.rx or args.require_ack:
|
||||
receiver = Receiver(ser, verbose=args.rx)
|
||||
receiver.start()
|
||||
|
||||
print(f"--> TX {cmd_name}: {packet.hex(' ')}")
|
||||
ser.write(packet)
|
||||
ser.flush()
|
||||
|
||||
if args.require_ack:
|
||||
if not receiver.ack_event.wait(args.ack_timeout):
|
||||
print(f"ERROR: no ACK within {args.ack_timeout}s", file=sys.stderr)
|
||||
exit_code = EXIT_NO_ACK
|
||||
elif receiver.result == "NACK":
|
||||
print("ERROR: device replied NACK", file=sys.stderr)
|
||||
exit_code = EXIT_NACK
|
||||
else:
|
||||
print("ACK received.")
|
||||
|
||||
if args.rx:
|
||||
time.sleep(args.rx_time)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
if receiver:
|
||||
receiver.stop_event.set()
|
||||
receiver.join(timeout=1)
|
||||
ser.close()
|
||||
|
||||
print("Done.")
|
||||
return exit_code
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user