dotfiles/.config/waybar/network.py

#!/usr/bin/env python3
"""Waybar network module: shows interface, IP, and up/down transfer rate."""

import argparse
import json
import os
import re
import subprocess
import time

STATE_FILE_TMPL = "/tmp/waybar-network-stats-{}"
ICON = "󰲝"  # nerd font network icon
NO_NET_ICON = "󰖪"
DEFAULT_ICON_SIZE = "24pt"


def run(cmd: list[str]) -> str:
    """Run a command and return stdout, or empty string on failure."""
    try:
        return subprocess.check_output(
            cmd, stderr=subprocess.DEVNULL, text=True
        ).strip()
    except subprocess.CalledProcessError:
        return ""


def get_default_interface() -> str | None:
    """Return the default route interface name."""
    output = run(["ip", "route"])
    for line in output.splitlines():
        if line.startswith("default"):
            parts = line.split()
            # "default via X.X.X.X dev INTERFACE ..."
            if len(parts) >= 5:
                return parts[4]
    return None


def get_ip(iface: str) -> str | None:
    """Return the IPv4 address of the given interface."""
    output = run(["ip", "-4", "addr", "show", iface])
    m = re.search(r"inet\s+(\d+\.\d+\.\d+\.\d+)", output)
    return m.group(1) if m else None


def iface_type(iface: str) -> str:
    """Classify interface as WiFi, Ethernet, or generic."""
    if re.match(r"^(wlan|wlp)", iface):
        return "WiFi"
    if re.match(r"^(eth|enp|eno|ens)", iface):
        return "Ethernet"
    return iface


def format_rate(bps: float) -> str:
    """Format bytes-per-second as human-readable with one decimal."""
    bps = max(bps, 0)
    if bps >= 1_073_741_824:
        return f"{bps / 1_073_741_824:.1f}G"
    if bps >= 1_048_576:
        return f"{bps / 1_048_576:.1f}M"
    if bps >= 1024:
        return f"{bps / 1024:.1f}K"
    return f"{bps:.1f}B"


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--icon-size", default=DEFAULT_ICON_SIZE,
                    help="Icon size (Pango size, e.g. 24pt)")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    isize = args.icon_size
    iface = get_default_interface()
    if iface is None:
        print(json.dumps(
            {"text": f"{NO_NET_ICON}  No net",
             "tooltip": "No network connection"},
            ensure_ascii=False))
        return

    ip = get_ip(iface)
    if ip is None:
        print(json.dumps(
            {"text": f"{NO_NET_ICON}  No IP",
             "tooltip": f"Interface {iface} has no IP"},
            ensure_ascii=False))
        return

    itype = iface_type(iface)

    # Read current byte counters
    now_ts = time.time_ns()
    rx_path = f"/sys/class/net/{iface}/statistics/rx_bytes"
    tx_path = f"/sys/class/net/{iface}/statistics/tx_bytes"

    try:
        with open(rx_path) as f:
            now_rx = int(f.read().strip())
        with open(tx_path) as f:
            now_tx = int(f.read().strip())
    except (OSError, ValueError):
        now_rx = now_tx = 0

    # Transfer rate calculation
    rate_down = ""
    rate_up = ""
    state_file = STATE_FILE_TMPL.format(iface)

    if os.path.exists(state_file):
        try:
            with open(state_file) as f:
                prev_rx, prev_tx, prev_ts = map(int, f.read().split())
            delta_ns = now_ts - prev_ts
            if delta_ns > 0:
                delta_rx = now_rx - prev_rx
                delta_tx = now_tx - prev_tx
                rx_bps = delta_rx * 1e9 / delta_ns
                tx_bps = delta_tx * 1e9 / delta_ns
                rate_down = format_rate(rx_bps)
                rate_up = format_rate(tx_bps)
        except (OSError, ValueError):
            pass

    # Save current state for next run
    with open(state_file, "w") as f:
        f.write(f"{now_rx} {now_tx} {now_ts}")

    # Build output
    if rate_down and rate_up:
        rate_str = f"  ↓{rate_down} ↑{rate_up}"
        tooltip_rate = f"↓ {rate_down}/s  ↑ {rate_up}/s"
    else:
        rate_str = ""
        tooltip_rate = "↓ ?/s  ↑ ?/s"

    text = (
        f"<span size='{isize}' rise='-1500'>{ICON}</span>"
        f"  {iface} {ip}{rate_str}"
    )
    tooltip = f"{itype}: {iface} — {ip} | {tooltip_rate}"

    print(json.dumps({"text": text, "tooltip": tooltip}, ensure_ascii=False))


if __name__ == "__main__":
    main()