ipyrf 0.1.0__py3-none-any.whl

ipyrf/__init__.py

@@ -0,0 +1,5 @@
+"""ipyrf package: a tiny iperf3-like tool with JSON output."""
+
+__all__ = [
+    "main",
+]

ipyrf/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    main()

ipyrf/cli.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+import argparse
+import sys
+
+from .logger import Logger
+from .utils import parse_bandwidth, parse_ip, tcp_congestion_control_info
+from . import tcp, udp
+from .interactive import InteractiveController
+from .controllers import StaticPacingController
+
+
+def main():
+    p = argparse.ArgumentParser(description="Minimal iperf3-like tool (JSON output)")
+
+    common = argparse.ArgumentParser(add_help=False)
+    common.add_argument("--port", type=int, default=5201, help="Port number")
+    common.add_argument(
+        "--logfile", help="Write log messages to a file instead of stdout"
+    )
+    common.add_argument(
+        "--json_log",
+        help="Write the log messages as JSON.",
+        action="store_true",
+        default=False,
+    )
+    common.add_argument("--interval", type=float, default=1.0, help="Stats interval")
+
+    subp = p.add_subparsers(dest="protocol", required=True)
+
+    tcp_parser = subp.add_parser("tcp", help="TCP mode")
+    tcp_sub = tcp_parser.add_subparsers(dest="role", required=True)
+
+    congestion_control = tcp_congestion_control_info()
+    common_tcp = argparse.ArgumentParser(add_help=False)
+    common_tcp.add_argument(
+        "--congestion-control",
+        choices=congestion_control.get("allowed", []),
+        default=None,
+        help=(
+            argparse.SUPPRESS
+            if congestion_control == {}
+            else (
+                "TCP: set congestion control algorithm "
+                f"(default: system default '{congestion_control.get('current')}')"
+            )
+        ),
+    )
+
+    tcp_srv = tcp_sub.add_parser(
+        "server", parents=[common, common_tcp], help="Run a TCP server"
+    )
+    tcp_srv.add_argument(
+        "address", metavar="ADDRESS", type=parse_ip, help="Listen address"
+    )
+
+    tcp_cli = tcp_sub.add_parser(
+        "client", parents=[common, common_tcp], help="Run a TCP client"
+    )
+    tcp_cli.add_argument(
+        "--bandwidth", type=parse_bandwidth, help="Target bandwidth, e.g., 50M"
+    )
+    tcp_cli.add_argument("address", metavar="ADDRESS", help="Server address to connect")
+    tcp_cli.add_argument(
+        "--set-mss", dest="set_mss", type=int, help="TCP: set TCP_MAXSEG (approx MSS)"
+    )
+
+    # Time and interactive mode are mutually exclusive
+    tcp_time_group = tcp_cli.add_mutually_exclusive_group()
+    tcp_time_group.add_argument(
+        "--time", type=int, default=10, help="Test duration in seconds"
+    )
+    tcp_time_group.add_argument(
+        "--interactive", action="store_true", help="Run client in interactive mode"
+    )
+
+    udp_parser = subp.add_parser("udp", help="UDP mode")
+    udp_sub = udp_parser.add_subparsers(dest="role", required=True)
+
+    udp_srv = udp_sub.add_parser("server", parents=[common], help="Run a UDP server")
+    udp_srv.add_argument(
+        "address", metavar="ADDRESS", type=parse_ip, help="Listen address"
+    )
+
+    udp_cli = udp_sub.add_parser("client", parents=[common], help="Run a UDP client")
+    udp_cli.add_argument("address", metavar="ADDRESS", help="Server address to connect")
+    udp_cli.add_argument(
+        "--bandwidth", type=parse_bandwidth, help="Target bandwidth, e.g., 50M"
+    )
+    udp_cli.add_argument(
+        "-l", dest="length", type=int, default=1200, help="UDP payload length"
+    )
+
+    # Time and interactive mode are mutually exclusive
+    udp_time_group = udp_cli.add_mutually_exclusive_group()
+    udp_time_group.add_argument(
+        "--time", type=int, default=10, help="Test duration in seconds"
+    )
+    udp_time_group.add_argument(
+        "--interactive", action="store_true", help="Run client in interactive mode"
+    )
+
+    args = p.parse_args()
+
+    if args.role not in ("server", "client"):
+        raise ValueError(f"Invalid role: {args.role}. Must be 'server' or 'client'.")
+
+    log = Logger(args.json_log, args.protocol, args.role, args.logfile)
+
+    controller = None
+    if args.protocol == "udp":
+        if args.role == "server":
+            udp.server(log, args.address, args.port, args.interval)
+        else:
+            bw = (
+                args.bandwidth or parse_bandwidth("50M")
+                if args.interactive
+                else (args.bandwidth or 1e9)
+            )
+            if args.interactive:
+                controller = InteractiveController(bw, args.length, args.interval)
+            else:
+                controller = StaticPacingController(
+                    bw, max(args.length, udp.UDP_HDR.size), args.time, args.interval
+                )
+            udp.client(
+                log,
+                args.address,
+                args.port,
+                args.length,
+                controller,
+            )
+
+    else:
+        if args.role == "server":
+            tcp.server(
+                log, args.address, args.port, args.interval, args.congestion_control
+            )
+        else:
+            if args.interactive:
+                quantum = args.set_mss if args.set_mss else 1200
+                # If no bandwidth provided, controller will act as unlimited until adjusted
+                controller = InteractiveController(
+                    args.bandwidth, quantum, args.interval
+                )
+            else:
+                controller = StaticPacingController(
+                    args.bandwidth,
+                    args.set_mss if args.set_mss else 1200,
+                    args.time,
+                    args.interval,
+                )
+            tcp.client(
+                log,
+                args.address,
+                args.port,
+                args.congestion_control,
+                args.set_mss,
+                controller,
+            )
+
+    if controller is not None:
+        controller.stop()

ipyrf/controllers.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+import time
+from typing import Dict, Optional
+
+from .token_bucket import TokenBucket
+
+
+class BasePacingController:
+    def __init__(self, interval_seconds: float = 1.0):
+        self.interval_seconds = interval_seconds
+
+    def is_pacing(self) -> bool:
+        return False
+
+    def maybe_sleep(self, n_bytes: int):
+        return
+
+    def get_update_fields(self) -> Dict[str, float]:
+        return {}
+
+    def should_stop(self) -> bool:
+        return False
+
+    def stop_reason(self) -> str:
+        return "unknown"
+
+    def start(self):
+        pass
+
+    def stop(self):
+        pass
+
+
+class StaticPacingController(BasePacingController):
+    def __init__(
+        self,
+        bandwidth_bps: Optional[float],
+        quantum_bytes: int,
+        duration_seconds: float,
+        interval_seconds: float,
+    ):
+        super().__init__(interval_seconds=interval_seconds)
+        self.bandwidth_bps = bandwidth_bps
+        self.duration_seconds = duration_seconds
+        self.start_time = None
+        self.tb: Optional[TokenBucket] = None
+        if bandwidth_bps is not None:
+            self.tb = TokenBucket(bandwidth_bps, quantum_bytes)
+
+    def is_pacing(self) -> bool:
+        return self.tb is not None
+
+    def maybe_sleep(self, n_bytes: int):
+        if self.tb is None:
+            return
+        while True:
+            sleep_time = self.tb.take(n_bytes)
+            if sleep_time <= 0:
+                break
+            time.sleep(sleep_time)
+
+    def get_update_fields(self) -> Dict[str, float]:
+        if self.bandwidth_bps is None:
+            return {}
+        return {"target_bandwidth_bps": float(self.bandwidth_bps)}
+
+    def start(self):
+        """Call this when the test starts to begin duration tracking."""
+        self.start_time = time.time()
+
+    def should_stop(self) -> bool:
+        """Check if the test should stop based on duration."""
+        assert self.start_time is not None
+        return (time.time() - self.start_time) >= self.duration_seconds
+
+    def stop_reason(self) -> str:
+        return "duration"

ipyrf/interactive.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+import os
+import sys
+import time
+import threading
+import termios
+import fcntl
+import select
+
+from .utils import human_bps
+from .token_bucket import TokenBucket
+from .controllers import BasePacingController
+
+
+class DynamicTokenBucket(TokenBucket):
+    """Token bucket that allows updating the target rate at runtime."""
+
+    def set_rate_bps(self, new_rate_bps: float, quantum_bytes: int):
+        # Convert bits/sec -> bytes/sec as used internally by TokenBucket
+        self.rate_bps = max(1e-6, new_rate_bps / 8.0)
+        self.capacity = max(quantum_bytes * 2, int(self.rate_bps))
+        self.tokens = min(self.tokens, self.capacity)
+
+
+class KeyReader:
+    """Non-blocking key reader that recognizes arrow keys and a few chars.
+
+    Reads from /dev/tty (the controlling terminal), so stdout can be redirected.
+    Prevents terminal echo so arrow keys won't shift the cursor (no stray spaces).
+    """
+
+    def __init__(self):
+        self.tty_path = "/dev/tty"
+        self.fd = None
+        self.orig_attrs = None
+        self.orig_flags = None
+
+    def __enter__(self):
+        try:
+            self.fd = os.open(self.tty_path, os.O_RDONLY | os.O_NONBLOCK)
+        except OSError:
+            self.fd = sys.stdin.fileno()
+
+        self.orig_attrs = termios.tcgetattr(self.fd)
+        attrs = termios.tcgetattr(self.fd)
+        lflag = attrs[3]
+        attrs[3] = lflag & ~(termios.ECHO | termios.ICANON | termios.IEXTEN)
+        cc = list(attrs[6])
+        cc[termios.VMIN] = 0
+        cc[termios.VTIME] = 0
+        attrs[6] = cc
+        termios.tcsetattr(self.fd, termios.TCSADRAIN, attrs)
+
+        self.orig_flags = fcntl.fcntl(self.fd, fcntl.F_GETFL)
+        fcntl.fcntl(self.fd, fcntl.F_SETFL, self.orig_flags | os.O_NONBLOCK)
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        try:
+            if self.orig_attrs is not None:
+                termios.tcsetattr(self.fd, termios.TCSADRAIN, self.orig_attrs)
+            if self.orig_flags is not None:
+                fcntl.fcntl(self.fd, fcntl.F_SETFL, self.orig_flags)
+            if self.fd not in (None, sys.stdin.fileno()):
+                os.close(self.fd)
+        except Exception:
+            pass
+
+    def read_keys(self, stop_event: threading.Event, on_key):
+        buf = bytearray()
+        allowed_chars = {"q", "Q", "0", "u", "U"}
+        while not stop_event.is_set():
+            try:
+                rlist, _, _ = select.select([self.fd], [], [], 0.05)
+            except Exception:
+                rlist = []
+            if not rlist:
+                continue
+            try:
+                chunk = os.read(self.fd, 16)
+            except BlockingIOError:
+                continue
+            except Exception:
+                break
+            if not chunk:
+                continue
+            buf.extend(chunk)
+            while True:
+                if not buf:
+                    break
+                c = buf[0]
+                if c == 0x1B:  # ESC
+                    if len(buf) >= 3 and buf[1] == 0x5B:
+                        code = buf[2]
+                        if code in (0x41, 0x42, 0x43, 0x44):  # A,B,C,D
+                            del buf[:3]
+                            mapping = {
+                                0x41: "UP",
+                                0x42: "DOWN",
+                                0x43: "RIGHT",
+                                0x44: "LEFT",
+                            }
+                            on_key(mapping[code])
+                            continue
+                    if len(buf) < 3:
+                        break
+                    del buf[0]
+                    continue
+                else:
+                    ch = chr(c)
+                    del buf[0]
+                    if ch in allowed_chars:
+                        on_key(ch)
+                    continue
+
+
+class InteractiveController(BasePacingController):
+    def __init__(
+        self, initial_bps: float | None, quantum_bytes: int, interval: float = 1.0
+    ):
+        super().__init__(interval_seconds=interval)
+        self.lock = threading.Lock()
+        self.stop_event = threading.Event()
+        self.quantum = quantum_bytes
+        self.pacing = initial_bps is not None
+        self.target_bps = (
+            float(initial_bps) if initial_bps is not None else float("inf")
+        )
+        self.tb: DynamicTokenBucket | None = (
+            DynamicTokenBucket(self.target_bps, self.quantum) if self.pacing else None
+        )
+        self.keyloop_thread = start_interactive_keys(self, initial_bps)
+
+    def is_pacing(self) -> bool:
+        with self.lock:
+            return self.pacing and self.target_bps != float("inf")
+
+    def maybe_sleep(self, n_bytes: int):
+        tb = None
+        with self.lock:
+            tb = self.tb
+        if tb is None:
+            return
+        while True:
+            sleep_time = tb.take(n_bytes)
+            if sleep_time <= 0:
+                break
+            time.sleep(sleep_time)
+
+    def get_update_fields(self):
+        with self.lock:
+            if self.pacing and self.target_bps != float("inf"):
+                return {"target_bandwidth_bps": self.target_bps}
+            return {}
+
+    def should_stop(self) -> bool:
+        return self.stop_event.is_set()
+
+    def stop_reason(self) -> str:
+        return "user-stop"
+
+    def reset(self, initial_bps: float | None):
+        with self.lock:
+            if initial_bps is None:
+                self.target_bps = float("inf")
+                self.pacing = False
+                self.tb = None
+            else:
+                self.target_bps = float(initial_bps)
+                self.pacing = True
+                if self.tb is None:
+                    self.tb = DynamicTokenBucket(self.target_bps, self.quantum)
+                else:
+                    self.tb.set_rate_bps(self.target_bps, self.quantum)
+
+    def unlimited(self):
+        with self.lock:
+            self.pacing = False
+            self.target_bps = float("inf")
+            self.tb = None
+
+    def bump(self, delta_bps: float = 0.0, scale: float = 1.0):
+        with self.lock:
+            if not self.pacing or self.target_bps == float("inf"):
+                self.pacing = True
+                if not (self.target_bps != float("inf")):
+                    self.target_bps = 50e6
+                if self.tb is None:
+                    self.tb = DynamicTokenBucket(self.target_bps, self.quantum)
+            if scale != 1.0:
+                self.target_bps = max(1e3, self.target_bps * scale)
+            else:
+                self.target_bps = max(1e3, self.target_bps + delta_bps)
+            self.tb.set_rate_bps(self.target_bps, self.quantum)
+            return self.target_bps
+
+    def request_stop(self):
+        self.stop_event.set()
+
+    def stop(self):
+        if self.keyloop_thread is not None:
+            self.keyloop_thread.join()
+            self.keyloop_thread = None
+
+
+def start_interactive_keys(
+    controller: InteractiveController, initial_bps: float | None
+):
+    def on_key(k: str):
+        if k == "RIGHT":
+            new = controller.bump(delta_bps=1e6)
+            print(f"[interactive] target = {human_bps(new)}", file=sys.stderr)
+        elif k == "LEFT":
+            new = controller.bump(delta_bps=-1e6)
+            print(f"[interactive] target = {human_bps(new)}", file=sys.stderr)
+        elif k == "UP":
+            new = controller.bump(scale=1.10)
+            print(f"[interactive] target = {human_bps(new)}", file=sys.stderr)
+        elif k == "DOWN":
+            new = controller.bump(scale=0.90)
+            print(f"[interactive] target = {human_bps(new)}", file=sys.stderr)
+        elif k == "0":
+            controller.reset(initial_bps)
+            print(
+                f"[interactive] reset -> {human_bps(initial_bps if initial_bps is not None else float('inf'))}",
+                file=sys.stderr,
+            )
+        elif k in ("u", "U"):
+            controller.unlimited()
+            print("[interactive] pacing: unlimited", file=sys.stderr)
+        elif k in ("q", "Q"):
+            controller.request_stop()
+
+    print(
+        "[interactive] Controls: ← -1 Mbps, -> +1 Mbps, ↓ -10%, ↑ +10%, 0 reset, u unlimited, q quit, Ctrl+C exit",
+        file=sys.stderr,
+    )
+    print(
+        f"[interactive] starting at {human_bps(initial_bps if initial_bps is not None else float('inf'))}",
+        file=sys.stderr,
+    )
+
+    def _keyloop():
+        with KeyReader() as kr:
+            kr.read_keys(controller.stop_event, on_key)
+
+    keyloop_thread = threading.Thread(target=_keyloop, daemon=True)
+    keyloop_thread.start()
+    return keyloop_thread

ipyrf/logger.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+import json
+import sys
+
+from .utils import human_readable_bytes
+
+
+log_types = ["start", "test", "update", "summary"]
+log_directions = ["tx", "rx"]
+log_modes = ["tcp", "udp"]
+
+
+class Logger:
+    def __init__(
+        self, json_log: bool, mode: str, role: str, logfile: str | None = None
+    ):
+        self.json_log = json_log
+        assert mode in log_modes
+        self.mode = mode
+        self.direction = "tx" if role == "client" else "rx"
+        assert self.direction in log_directions
+        self.test_start_time = None
+        self.logfile = logfile
+
+    def start(self, ip: str, port: int):
+        self._log(
+            log_type="start",
+            mode=self.mode,
+            direction=self.direction,
+            address=f"{ip}:{port}",
+        )
+
+    def test(self, peer_ip: str, peer_port: int, start_ts: float):
+        self._log(
+            log_type="test",
+            peer=f"{peer_ip}:{peer_port}",
+            ts=start_ts,
+        )
+        self.test_start_time = start_ts
+
+    def update(self, start_ts: float, end_ts: float, bytes: int, **obj):
+        if self.test_start_time is None:
+            raise RuntimeError("test() must be called before update()")
+        end_ts -= self.test_start_time
+        start_ts -= self.test_start_time
+        delta_t = end_ts - start_ts
+
+        bps = (bytes * 8.0) / delta_t if delta_t > 0 else 0.0
+        self._log(
+            log_type="update",
+            start=start_ts,
+            end=end_ts,
+            bytes=bytes,
+            bits_per_second=bps,
+            **obj,
+        )
+
+    def summary(self, **obj):
+        self._log(log_type="summary", **obj)
+
+    def write(self, message: str):
+        if self.logfile:
+            with open(self.logfile, "a", buffering=1) as f:
+                f.write(message + "\n")
+        else:
+            print(message)
+            sys.stdout.flush()
+
+    def _log(self, log_type, **obj):
+        if self.json_log:
+            obj["type"] = log_type
+            obj["mode"] = self.mode
+            obj["direction"] = self.direction
+            self.write(json.dumps(obj, separators=(",", ":")) + "\n")
+            return
+        assert log_type in log_types
+        if log_type == "start":
+            self.write(
+                f"▶ {self.mode.upper()} {self.direction.upper()} — {obj['address']}"
+            )
+        elif log_type == "test":
+            self.write(f"▶ TEST peer={obj['peer']}  ts={obj['ts']}")
+        elif log_type == "update":
+            message = (
+                f"⏱ {obj['start']:.2f}-{obj['end']:.2f} sec"
+                f" | {human_readable_bytes(obj['bytes'])}"
+            )
+            if "target_bandwidth_bps" in obj:
+                message += f" | {obj['bits_per_second'] / 1e6:.2f}/{obj['target_bandwidth_bps'] / 1e6:.2f} Mbps"
+            else:
+                message += f" | {obj['bits_per_second'] / 1e6:.2f} Mbps"
+            if "lost_packets" in obj and "packets" in obj:
+                message += f" | {obj['lost_packets']}/{obj['packets']} lost ({obj['lost_percent']:.1f}%)"
+            elif "packets" in obj:
+                message += f" | {obj['packets']} pkts"
+            self.write(message)
+        elif log_type == "summary":
+            self.write(
+                "\n━ SUMMARY ━\n"
+                f"  {self.mode.upper()} {self.direction.upper()}\n"
+                f"  {obj.get('sender', '')} → {obj.get('receiver', '')}\n"
+                f"  duration : {obj['seconds']:.2f} sec\n"
+                f"  data     : {human_readable_bytes(obj['bytes'])}\n"
+                f"  rate     : {obj['bits_per_second'] / 1e6:.2f} Mbps\n"
+                f"  reason   : {obj.get('stop_reason', '')}\n"
+            )