#!/usr/bin/env python3
"""retrans - per-connection TCP retransmit summary from `ss -tin`.

Usage: retrans [-n N] [--watch SECONDS] [--min-bytes B] [ss filter ...]

Examples:
    retrans                       # all established TCP connections
    retrans dport = :443          # only connections to port 443
    retrans --watch 2             # print deltas every 2 seconds
"""
import argparse, re, subprocess, sys, time

__version__ = "0.2.1"
FIELDS = re.compile(r"(bytes_sent|bytes_retrans|bytes_acked|rtt|cwnd|retrans):([\d./]+)")


def snapshot(filt):
    out = subprocess.run(["ss", "-tinH", "state", "established"] + filt,
                         capture_output=True, text=True, check=True).stdout
    conns, cur = {}, None
    for line in out.splitlines():
        if not line.startswith((" ", "\t")):
            parts = line.split()
            if len(parts) >= 4:
                cur = f"{parts[-2]} -> {parts[-1]}"
                conns[cur] = {}
            continue
        if cur:
            for k, v in FIELDS.findall(line):
                conns[cur][k] = v
    rows = {}
    for key, f in conns.items():
        sent = int(f.get("bytes_sent", 0))
        rtr = int(f.get("bytes_retrans", 0))
        rtt = float(f.get("rtt", "0/0").split("/")[0])
        rows[key] = (sent, rtr, rtt, int(f.get("cwnd", 0)))
    return rows


def show(rows, n, min_bytes):
    rows = [(k, *v) for k, v in rows.items() if v[0] >= min_bytes]
    rows.sort(key=lambda r: (r[2] / r[1]) if r[1] else 0, reverse=True)
    print(f"{'connection':<52} {'sent':>10} {'retrans':>9} {'ratio':>7} {'rtt ms':>7} {'cwnd':>5}")
    for k, sent, rtr, rtt, cwnd in rows[:n]:
        ratio = 100 * rtr / sent if sent else 0
        print(f"{k:<52} {human(sent):>10} {human(rtr):>9} {ratio:>6.2f}% {rtt:>7.1f} {cwnd:>5}")


def human(n):
    for u in ("B", "K", "M", "G"):
        if n < 1024:
            return f"{n:.0f}{u}"
        n /= 1024
    return f"{n:.1f}T"


def main(argv=None):
    ap = argparse.ArgumentParser(prog="retrans", description=__doc__.split("\n")[0])
    ap.add_argument("filter", nargs="*", help="ss filter expression")
    ap.add_argument("-n", type=int, default=20, help="rows to show (default 20)")
    ap.add_argument("--watch", type=float, help="print deltas every N seconds")
    ap.add_argument("--min-bytes", type=int, default=64 * 1024)
    ap.add_argument("--version", action="version", version=__version__)
    a = ap.parse_args(argv)

    if not a.watch:
        show(snapshot(a.filter), a.n, a.min_bytes)
        return 0
    prev = snapshot(a.filter)
    try:
        while True:
            time.sleep(a.watch)
            cur = snapshot(a.filter)
            delta = {k: (v[0] - prev[k][0], v[1] - prev[k][1], v[2], v[3]) for k, v in cur.items() if k in prev}
            print(time.strftime("\n%H:%M:%S"))
            show(delta, a.n, 1)
            prev = cur
    except KeyboardInterrupt:
        return 0


if __name__ == "__main__":
    sys.exit(main())
