test.py

#!/usr/bin/env python3
import argparse
import socket
import struct
import threading
import time
import sys
import signal

FRAME_SIZE = 13

def hex_bytes(s: str) -> bytes:
    """Parse a string like '12 34 56 78 00' or '1234567800' into bytes."""
    s = s.replace("0x", "").replace(" ", "").replace("_", "").replace(",", "")
    if len(s) % 2 != 0:
        raise ValueError("hex string must have even length")
    return bytes.fromhex(s)

def pack_frame(can_id: int, data: bytes, *, extended: bool, rtr: bool=False, dlc: int=None) -> bytes:
    """Pack to 13-byte wire format:
       [0]=flags+dlc, [1..4]=BE ID, [5..12]=8 data bytes (zero-padded)."""
    if dlc is None:
        dlc = len(data)
    if not (0 <= dlc <= 8):
        raise ValueError("DLC must be 0..8")
    if extended:
        if not (0 <= can_id <= 0x1FFFFFFF):
            raise ValueError("Extended ID must be 0..0x1FFFFFFF")
        id_be = can_id & 0x1FFFFFFF
    else:
        if not (0 <= can_id <= 0x7FF):
            raise ValueError("Standard ID must be 0..0x7FF")
        id_be = can_id & 0x7FF

    hdr = 0
    if extended: hdr |= 0x80
    if rtr:      hdr |= 0x40
    hdr |= (dlc & 0x0F)

    out = bytearray(FRAME_SIZE)
    out[0] = hdr
    out[1:5] = struct.pack(">I", id_be)
    data = data[:8].ljust(8, b"\x00")
    out[5:13] = data
    return bytes(out)

def unpack_frame(buf: bytes):
    """Return dict with ext, rtr, dlc, id, data (bytes len 8)."""
    if len(buf) != FRAME_SIZE:
        raise ValueError("bad frame size")
    hdr = buf[0]
    ext = bool(hdr & 0x80)
    rtr = bool(hdr & 0x40)
    dlc = hdr & 0x0F
    cid, = struct.unpack(">I", buf[1:5])
    if ext:
        can_id = cid & 0x1FFFFFFF
    else:
        can_id = cid & 0x7FF
    data = buf[5:13]
    return {"ext": ext, "rtr": rtr, "dlc": dlc, "id": can_id, "data": data}

def fmt_hex(b: bytes) -> str:
    return " ".join(f"{x:02X}" for x in b)

def rx_loop(name: str, sock: socket.socket, stop_event: threading.Event):
    buf = bytearray()
    try:
        while not stop_event.is_set():
            chunk = sock.recv(4096)
            if not chunk:
                print(f"[{name} RX] connection closed by server")
                break
            buf += chunk
            while len(buf) >= FRAME_SIZE:
                frame = bytes(buf[:FRAME_SIZE])
                del buf[:FRAME_SIZE]
                try:
                    parsed = unpack_frame(frame)
                    kind = "EXT" if parsed["ext"] else "STD"
                    print(f"[{name} RX] {fmt_hex(frame)}  | {kind} DLC={parsed['dlc']} ID=0x{parsed['id']:X} DATA={fmt_hex(parsed['data'][:parsed['dlc']])}")
                except Exception as e:
                    print(f"[{name} RX] decode error: {e}  raw={fmt_hex(frame)}")
    except Exception as e:
        print(f"[{name} RX] error: {e}")

def connect_client(name: str, host: str, port: int, stop_event: threading.Event):
    """Connect a single client and start its RX thread. Returns (socket, rx_thread)."""
    s = socket.create_connection((host, port))
    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    print(f"[{name}] connected to {host}:{port}")
    t = threading.Thread(target=rx_loop, args=(name, s, stop_event), daemon=True)
    t.start()
    return s, t

def main():
    ap = argparse.ArgumentParser(description="TCP test client for CAN↔TCP gateway (13-byte frames). Supports multiple clients.")
    ap.add_argument("--host", default="127.0.0.1", help="server host (default: 127.0.0.1)")
    ap.add_argument("--port", type=int, default=10001, help="server port (default: 10001)")
    ap.add_argument("--clients", type=int, default=1, help="number of clients to connect (default: 1)")

    g = ap.add_mutually_exclusive_group()
    g.add_argument("--frame-hex", help="full 13-byte frame hex (e.g. '85 12 34 56 78 12 34 56 78 00 00 00 00')")
    g.add_argument("--compose", action="store_true",
                   help="compose a frame from --ext/--std --id --data instead of --frame-hex")

    ap.add_argument("--ext", action="store_true", help="extended frame (for --compose)")
    ap.add_argument("--std", action="store_true", help="standard frame (for --compose)")
    ap.add_argument("--rtr", action="store_true", help="RTR frame (rare; for --compose)")
    ap.add_argument("--id",  help="CAN ID (e.g., 0x12345678 or 0x678) for --compose")
    ap.add_argument("--data", default="", help="data bytes for --compose, e.g. '12 34 56 78 00' (max 8 bytes)")

    ap.add_argument("--repeat", type=int, default=1, help="send this many times (default: 1)")
    ap.add_argument("--interval", type=float, default=0.25, help="seconds between repeated sends (default: 0.25)")
    ap.add_argument("--send-from", type=int, default=0, help="which client sends (0-indexed, default: 0). -1 = all clients send.")
    ap.add_argument("--interactive", action="store_true",
                    help="after sending, enter an interactive mode to type more frames")

    args = ap.parse_args()

    # Build outbound frame
    out_frames = []

    if args.frame_hex:
        b = hex_bytes(args.frame_hex)
        if len(b) != FRAME_SIZE:
            ap.error("--frame-hex must be exactly 13 bytes")
        out_frames.append(b)
    elif args.compose:
        if not (args.ext ^ args.std):
            ap.error("must specify exactly one of --ext or --std with --compose")
        if not args.id:
            ap.error("--id is required with --compose")
        try:
            can_id = int(args.id, 0)
        except ValueError:
            ap.error("invalid --id (use e.g. 0x123 or 0x12345678)")
        data = hex_bytes(args.data) if args.data else b""
        f = pack_frame(can_id, data, extended=args.ext, rtr=args.rtr)
        out_frames.append(f)
    else:
        f = pack_frame(0x12345678, bytes.fromhex("12 34 56 78 00"), extended=True, rtr=False)
        print("[INFO] No frame specified; sending default EXT id=0x12345678 data=12 34 56 78 00")
        out_frames.append(f)

    stop_event = threading.Event()

    # Connect all clients
    clients = []
    for i in range(args.clients):
        name = f"C{i}"
        s, t = connect_client(name, args.host, args.port, stop_event)
        clients.append((name, s, t))

    time.sleep(0.1)

    # Determine which clients send
    if args.send_from == -1:
        senders = list(range(len(clients)))
    else:
        if args.send_from >= len(clients):
            print(f"ERR: --send-from {args.send_from} but only {len(clients)} client(s)")
            sys.exit(1)
        senders = [args.send_from]

    # Send frames
    for i in range(args.repeat):
        for idx in senders:
            name, s, _ = clients[idx]
            for f in out_frames:
                print(f"[{name} TX] {fmt_hex(f)}")
                s.sendall(f)
        if i + 1 < args.repeat:
            time.sleep(args.interval)

    # Interactive mode
    if args.interactive:
        print(f"\nInteractive mode ({len(clients)} client(s) connected).")
        print("Prefix with client number to send from that client (default C0):")
        print("  0 ext: 0x12345678 12 34 56 78 00")
        print("  1 std: 0x678 12 34 56 78 00")
        print("  hex: 85 12 34 56 78 12 34 56 78 00 00 00 00")
        print("Ctrl+D or empty line to exit.\n")
        try:
            while True:
                line = input("> ").strip()
                if not line:
                    break
                try:
                    # Check for client prefix (digit at start)
                    client_idx = 0
                    parts = line.split(None, 1)
                    if parts[0].isdigit():
                        client_idx = int(parts[0])
                        if client_idx >= len(clients):
                            print(f"ERR: no client {client_idx}")
                            continue
                        line = parts[1] if len(parts) > 1 else ""

                    if line.lower().startswith("hex:"):
                        f = hex_bytes(line.split(":",1)[1])
                        if len(f) != FRAME_SIZE:
                            print("ERR: hex must be 13 bytes")
                            continue
                    else:
                        parts = line.split()
                        if len(parts) < 2:
                            print("ERR: need kind and id at least")
                            continue
                        kind = parts[0].lower()
                        can_id = int(parts[1], 0)
                        data = b""
                        if len(parts) > 2:
                            data = hex_bytes("".join(parts[2:]))
                        if kind == "ext":
                            f = pack_frame(can_id, data, extended=True)
                        elif kind == "std":
                            f = pack_frame(can_id, data, extended=False)
                        else:
                            print("ERR: kind must be 'ext' or 'std' or 'hex:'")
                            continue

                    name = clients[client_idx][0]
                    print(f"[{name} TX] {fmt_hex(f)}")
                    clients[client_idx][1].sendall(f)
                except Exception as e:
                    print("ERR:", e)
        except (EOFError, KeyboardInterrupt):
            pass

    # Graceful shutdown
    stop_event.set()
    for name, s, t in clients:
        try:
            s.shutdown(socket.SHUT_RDWR)
        except Exception:
            pass
        s.close()
        t.join(timeout=0.5)

if __name__ == "__main__":
    signal.signal(signal.SIGINT, lambda *args: sys.exit(0))
    main()