ping_collect.py

#!/usr/bin/env python3
import argparse
import csv
import json
import os
import re
import subprocess
import sys
import time
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple


TIME_RE = re.compile(r"time[=<]([0-9.]+)\s*ms")


@dataclass
class Location:
    name: str
    address: str


def build_ping_command(address: str, timeout_s: float) -> List[str]:
    if sys.platform.startswith("linux"):
        return ["ping", "-c", "1", "-W", str(int(timeout_s)), address]
    if sys.platform.startswith("darwin"):
        return ["ping", "-c", "1", "-W", str(int(timeout_s * 1000)), address]
    return ["ping", "-n", "1", "-w", str(int(timeout_s * 1000)), address]


def parse_latency_ms(output: str) -> Optional[float]:
    match = TIME_RE.search(output)
    if not match:
        return None
    return float(match.group(1))


def ping_once(address: str, timeout_s: float) -> Optional[float]:
    cmd = build_ping_command(address, timeout_s)
    try:
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            check=False,
            text=True,
        )
    except FileNotFoundError as exc:
        raise RuntimeError("ping command not found on PATH") from exc
    return parse_latency_ms(result.stdout)


def gather_latencies(
    locations: List[Location],
    duration_s: float,
    timeout_s: float,
) -> Dict[str, List[Tuple[float, float]]]:
    latencies: Dict[str, List[Tuple[float, float]]] = {loc.address: [] for loc in locations}
    start_time = time.monotonic()
    end_time = start_time + duration_s
    while time.monotonic() < end_time:
        for loc in locations:
            if time.monotonic() >= end_time:
                break
            latency = ping_once(loc.address, timeout_s)
            if latency is not None:
                elapsed = time.monotonic() - start_time
                latencies[loc.address].append((elapsed, latency))
                print(f"{loc.address} t={elapsed:.2f}s latency={latency:.2f}ms")
        if time.monotonic() < end_time:
            time.sleep(1.0)
    return latencies


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Ping a set of locations and save latency samples.")
    parser.add_argument(
        "--addresses-file",
        dest="addresses_file",
        default="addresses.json",
        help="Path to a JSON file containing a list of locations",
    )
    parser.add_argument(
        "--duration",
        type=float,
        default=60.0,
        help="Duration to ping in seconds",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=1.0,
        help="Ping timeout in seconds",
    )
    parser.add_argument(
        "--output-dir",
        default="ping_output",
        help="Output directory for data files",
    )
    return parser


def parse_args() -> argparse.Namespace:
    return build_parser().parse_args()


def load_locations(addresses_file: str) -> List[Location]:
    with open(addresses_file, "r", encoding="utf-8") as handle:
        data = handle.read()
    try:
        entries = json.loads(data)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Invalid JSON in {addresses_file}") from exc
    if not isinstance(entries, list):
        raise RuntimeError("addresses.json must contain a JSON array of objects")
    if not entries:
        raise RuntimeError("No addresses provided")
    locations: List[Location] = []
    for entry in entries:
        if not isinstance(entry, dict):
            raise RuntimeError("Each address entry must be a JSON object")
        name = entry.get("name")
        address = entry.get("address")
        if not isinstance(name, str) or not isinstance(address, str):
            raise RuntimeError("Each entry must include string fields: name and address")
        locations.append(Location(name=name, address=address))
    return locations


def save_json(
    output_path: str,
    locations: List[Location],
    latencies: Dict[str, List[Tuple[float, float]]],
    duration_s: float,
) -> None:
    location_payload = [{"name": loc.name, "address": loc.address} for loc in locations]
    samples = []
    for loc in locations:
        for elapsed, latency in latencies[loc.address]:
            samples.append(
                {
                    "name": loc.name,
                    "address": loc.address,
                    "elapsed_s": round(elapsed, 6),
                    "latency_ms": round(latency, 6),
                }
            )
    payload = {
        "duration_s": duration_s,
        "locations": location_payload,
        "samples": samples,
    }
    with open(output_path, "w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2)


def save_csv(
    output_path: str,
    locations: List[Location],
    latencies: Dict[str, List[Tuple[float, float]]],
) -> None:
    with open(output_path, "w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["name", "address", "elapsed_s", "latency_ms"])
        for loc in locations:
            for elapsed, latency in latencies[loc.address]:
                writer.writerow([loc.name, loc.address, f"{elapsed:.6f}", f"{latency:.6f}"])


def main() -> int:
    if len(sys.argv) == 1:
        build_parser().print_help()
        return 2
    args = parse_args()
    locations = load_locations(args.addresses_file)
    latencies = gather_latencies(locations, args.duration, args.timeout)
    os.makedirs(args.output_dir, exist_ok=True)
    json_path = os.path.join(args.output_dir, "ping_samples.json")
    csv_path = os.path.join(args.output_dir, "ping_samples.csv")
    save_json(json_path, locations, latencies, args.duration)
    save_csv(csv_path, locations, latencies)
    print(f"Saved data to {json_path} and {csv_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())