udpscope/tools/mock_tcpstreamer.py

#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: 2026 Martino Ferrari
"""A minimal MARTe2 TCPStreamer / TCPStreamService producer, for exercising udpscope.

Listens on `--port`, serves one client at a time with the same frames as
mock_udpstreamer.py (CONFIG, then Accumulate DATA packets of sine/square/ramp/noise), each
framed UDPS-over-TCP. Without `--stream` it behaves like a standalone TCPStreamer: the
stream starts at once and anything the client sends is ignored. With `--stream NAME` it
behaves like a TCPStreamService: the client must open with a SUBSCRIBE frame naming NAME
within 5 s, otherwise the connection is closed.

    ./tools/mock_tcpstreamer.py --port 44500 --stream fast --close-after 6
"""

import argparse
import os
import select
import socket
import struct
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mock_udpstreamer as m  # noqa: E402

PKT_SUBSCRIBE = 5
HEADER = 17


def read_exact(conn, n, deadline):
    buf = b""
    while len(buf) < n:
        left = deadline - time.monotonic()
        if left <= 0:
            return None
        r, _, _ = select.select([conn], [], [], left)
        if not r:
            return None
        chunk = conn.recv(n - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf


def subscribed_to(conn, want):
    """True once a SUBSCRIBE frame naming `want` arrived within 5 s."""
    deadline = time.monotonic() + 5.0
    head = read_exact(conn, HEADER, deadline)
    if head is None:
        return False
    magic, ptype, _, _, _, plen = struct.unpack("<IBIHHI", head)
    if magic != m.MAGIC or ptype != PKT_SUBSCRIBE or not 0 < plen <= 4096:
        return False
    name = read_exact(conn, plen, deadline)
    return name is not None and name.decode(errors="replace") == want


def serve(conn, args):
    dt = 1.0 / args.rate
    period = args.batch / args.rate
    conn.sendall(m.build_config(args.rate, m.TIME_FULL_ARRAY, "V"))
    counter, t_stream = 1, 0.0
    started = time.monotonic()
    next_send = started
    while True:
        now = time.monotonic()
        if args.close_after > 0 and now - started >= args.close_after:
            print("mock tcpstreamer: closing the connection (--close-after)", flush=True)
            return
        if now < next_send:
            time.sleep(min(next_send - now, 0.002))
            continue
        pkt = m.build_data(counter, t_stream, dt, args.batch, args.freq)
        if not (args.drop_every > 0 and counter % args.drop_every == 0):
            conn.sendall(pkt)
        counter += 1
        t_stream += args.batch * dt
        next_send += period


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=44500)
    ap.add_argument("--stream", default="", help="TCPStreamService mode: require SUBSCRIBE of this name")
    ap.add_argument("--rate", type=float, default=10000.0, help="samples per second")
    ap.add_argument("--batch", type=int, default=100, help="samples per packet")
    ap.add_argument("--freq", type=float, default=5.0, help="waveform frequency, Hz")
    ap.add_argument("--close-after", type=float, default=0.0, metavar="S",
                    help="close each connection after S seconds (the client must reconnect); 0 = never")
    ap.add_argument("--drop-every", type=int, default=0, metavar="N",
                    help="skip sending every Nth packet but advance the counter (a server-side drop); 0 = never")
    args = ap.parse_args()

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("0.0.0.0", args.port))
    srv.listen(4)
    print("mock tcpstreamer: listening on :%d (%s)" % (
        args.port, "service stream '%s'" % args.stream if args.stream else "standalone"), flush=True)
    while True:
        conn, addr = srv.accept()
        print("mock tcpstreamer: client %s:%d" % addr, flush=True)
        try:
            if args.stream and not subscribed_to(conn, args.stream):
                print("mock tcpstreamer: no valid SUBSCRIBE for '%s'; closing" % args.stream, flush=True)
                continue
            if args.stream:
                print("mock tcpstreamer: SUBSCRIBE '%s'" % args.stream, flush=True)
            serve(conn, args)
        except (BrokenPipeError, ConnectionResetError):
            print("mock tcpstreamer: client went away", flush=True)
        finally:
            conn.close()


if __name__ == "__main__":
    main()