FabledSteward/plugins/host_agent/agent.py

# plugins/host_agent/agent.py
"""Steward host agent — pushes resource metrics to a Steward instance.

Python 3.8+ stdlib only. Target ~300 lines. Served to targets at
GET /plugins/host_agent/agent.py.
"""
from __future__ import annotations

import json
import os
import shutil
import signal
import socket
import sys
import time
import urllib.error
import urllib.request
from collections import deque
from datetime import datetime, timezone

AGENT_VERSION = "1.1.0"


class ConfigError(Exception):
    pass


REQUIRED_KEYS = ("url", "token")
INT_KEYS = ("interval_seconds",)
LIST_KEYS = ("mounts",)


def read_config(path: str) -> dict:
    """Parse a flat `key = value` config file."""
    cfg: dict = {}
    try:
        with open(path, "r", encoding="utf-8") as f:
            for lineno, raw in enumerate(f, 1):
                line = raw.strip()
                if not line or line.startswith("#"):
                    continue
                if "=" not in line:
                    raise ConfigError(f"{path}:{lineno}: expected 'key = value'")
                key, _, value = line.partition("=")
                key = key.strip()
                value = value.strip()
                if key in INT_KEYS:
                    try:
                        cfg[key] = int(value)
                    except ValueError:
                        raise ConfigError(f"{path}:{lineno}: {key} must be int")
                elif key in LIST_KEYS:
                    cfg[key] = [v.strip() for v in value.split(",") if v.strip()]
                else:
                    cfg[key] = value
    except FileNotFoundError:
        raise ConfigError(f"{path}: not found")

    missing = [k for k in REQUIRED_KEYS if k not in cfg]
    if missing:
        raise ConfigError(f"{path}: missing required key(s): {', '.join(missing)}")

    cfg.setdefault("interval_seconds", 30)
    return cfg


# ─── collectors ──────────────────────────────────────────────────────────────

STAT_PATH = "/proc/stat"
MEMINFO_PATH = "/proc/meminfo"
LOADAVG_PATH = "/proc/loadavg"
UPTIME_PATH = "/proc/uptime"
OS_RELEASE_PATH = "/etc/os-release"


def _read_file(path: str) -> str:
    with open(path, "r", encoding="utf-8") as f:
        return f.read()


def _parse_cpu_line(stat_text: str) -> tuple[int, int]:
    """Return (total_jiffies, idle_jiffies) for the aggregate cpu line."""
    for line in stat_text.splitlines():
        if line.startswith("cpu "):
            parts = line.split()
            fields = [int(x) for x in parts[1:]]
            idle = fields[3] + (fields[4] if len(fields) > 4 else 0)  # idle + iowait
            total = sum(fields)
            return total, idle
    raise RuntimeError("no aggregate cpu line in /proc/stat")


def collect_cpu(sample_window: float = 0.2) -> float:
    """Return CPU utilization % over sample_window seconds."""
    t1, i1 = _parse_cpu_line(_read_file(STAT_PATH))
    time.sleep(sample_window)
    t2, i2 = _parse_cpu_line(_read_file(STAT_PATH))
    dt = t2 - t1
    di = i2 - i1
    if dt <= 0:
        return 0.0
    return round(100.0 * (dt - di) / dt, 2)


def collect_memory() -> dict:
    info: dict[str, int] = {}
    for line in _read_file(MEMINFO_PATH).splitlines():
        if ":" not in line:
            continue
        key, _, rest = line.partition(":")
        parts = rest.strip().split()
        if not parts:
            continue
        try:
            value_kb = int(parts[0])
        except ValueError:
            continue
        info[key] = value_kb * 1024  # bytes
    total = info.get("MemTotal", 0)
    available = info.get("MemAvailable", 0)
    swap_total = info.get("SwapTotal", 0)
    swap_free = info.get("SwapFree", 0)
    return {
        "total_bytes": total,
        "used_bytes": max(total - available, 0),
        "available_bytes": available,
        "swap_used_bytes": max(swap_total - swap_free, 0),
    }


def collect_storage(mounts: list[str]) -> list[dict]:
    out: list[dict] = []
    for m in mounts:
        try:
            u = shutil.disk_usage(m)
        except (FileNotFoundError, PermissionError):
            continue
        out.append({"mount": m, "total_bytes": u.total, "used_bytes": u.used})
    return out


def collect_load() -> dict:
    parts = _read_file(LOADAVG_PATH).split()
    return {"1m": float(parts[0]), "5m": float(parts[1]), "15m": float(parts[2])}


def collect_uptime() -> int:
    return int(float(_read_file(UPTIME_PATH).split()[0]))


def detect_primary_ip() -> str | None:
    """Best-effort primary non-loopback IPv4 of this host.

    connect() on a SOCK_DGRAM socket only sets the default peer and selects the
    outbound interface — no packets are sent — so getsockname() reveals the
    source IP the kernel would use to reach the internet. This is the host's own
    view of its address, so it survives reverse proxies / NAT (unlike the
    server reading request.remote_addr). Returns None on any error (offline / no
    route) and skips loopback.
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
    except OSError:
        return None
    finally:
        s.close()
    if not ip or ip.startswith("127."):
        return None
    return ip


def collect_metadata() -> dict:
    u = os.uname()
    distro = "unknown"
    try:
        for line in _read_file(OS_RELEASE_PATH).splitlines():
            if line.startswith("PRETTY_NAME="):
                distro = line.split("=", 1)[1].strip().strip('"')
                break
    except (OSError, FileNotFoundError):
        pass
    return {
        "kernel": u.release,
        "distro": distro,
        "arch": u.machine,
        "host_ip": detect_primary_ip(),
    }


def default_mounts() -> list[str]:
    """Return real mount points from /proc/mounts, skipping pseudo filesystems."""
    skip_types = {"tmpfs", "devtmpfs", "proc", "sysfs", "cgroup", "cgroup2",
                  "overlay", "squashfs", "ramfs", "devpts", "mqueue", "pstore",
                  "securityfs", "debugfs", "tracefs", "hugetlbfs", "fusectl",
                  "configfs", "bpf", "autofs", "efivarfs", "binfmt_misc"}
    mounts: list[str] = []
    try:
        with open("/proc/mounts", "r") as f:
            for line in f:
                parts = line.split()
                if len(parts) < 3:
                    continue
                if parts[2] in skip_types:
                    continue
                mounts.append(parts[1])
    except OSError:
        return ["/"]
    return mounts or ["/"]


# ─── ring buffer ─────────────────────────────────────────────────────────────


class RingBuffer:
    def __init__(self, maxlen: int = 20) -> None:
        self._dq: deque = deque(maxlen=maxlen)

    def push(self, item) -> None:
        self._dq.append(item)

    def drain(self) -> list:
        out = list(self._dq)
        self._dq.clear()
        return out

    def __len__(self) -> int:
        return len(self._dq)


# ─── payload ─────────────────────────────────────────────────────────────────


def build_sample(mounts: list[str]) -> dict:
    """Collect one full sample. Partial samples allowed if a collector fails."""
    sample: dict = {"ts": datetime.now(timezone.utc).isoformat()}
    try:
        sample["cpu_pct"] = collect_cpu()
    except Exception:
        sample["cpu_pct"] = None
    try:
        sample["mem"] = collect_memory()
    except Exception:
        sample["mem"] = None
    try:
        sample["load"] = collect_load()
    except Exception:
        sample["load"] = None
    try:
        sample["uptime_secs"] = collect_uptime()
    except Exception:
        sample["uptime_secs"] = None
    try:
        sample["storage"] = collect_storage(mounts)
    except Exception:
        sample["storage"] = []
    return sample


def build_payload(samples: list[dict], hostname: str, metadata: dict) -> dict:
    return {
        "agent_version": AGENT_VERSION,
        "hostname": hostname,
        "metadata": metadata,
        "samples": samples,
    }


# ─── POST + backoff ──────────────────────────────────────────────────────────


BACKOFF_CAP = 300


def next_backoff(current: int) -> int:
    if current <= 0:
        return 30
    return min(current * 2, BACKOFF_CAP)


def post_payload(url: str, token: str, payload: dict) -> tuple[bool, int | None]:
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url.rstrip("/") + "/plugins/host_agent/ingest",
        data=body,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}",
            "User-Agent": f"steward-host-agent/{AGENT_VERSION}",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            return (200 <= resp.status < 300, resp.status)
    except urllib.error.HTTPError as e:
        return (False, e.code)
    except (urllib.error.URLError, TimeoutError, socket.timeout, OSError):
        return (False, None)


# ─── main loop ───────────────────────────────────────────────────────────────


_reload_requested = False
_shutdown_requested = False


def _handle_hup(_signum, _frame):
    global _reload_requested
    _reload_requested = True


def _handle_term(_signum, _frame):
    global _shutdown_requested
    _shutdown_requested = True


def _log(level: str, msg: str) -> None:
    sys.stderr.write(f"[{level}] {msg}\n")
    sys.stderr.flush()


def main_loop(conf_path: str) -> int:
    global _reload_requested, _shutdown_requested
    signal.signal(signal.SIGHUP, _handle_hup)
    signal.signal(signal.SIGTERM, _handle_term)

    try:
        cfg = read_config(conf_path)
    except ConfigError as e:
        _log("ERROR", str(e))
        return 2

    metadata = collect_metadata()
    hostname = cfg.get("hostname") or socket.gethostname()
    mounts = cfg.get("mounts") or default_mounts()
    buffer = RingBuffer(maxlen=20)
    backoff = 0

    _log("INFO", f"steward-host-agent {AGENT_VERSION} starting "
                 f"(url={cfg['url']}, interval={cfg['interval_seconds']}s)")

    while not _shutdown_requested:
        if _reload_requested:
            try:
                cfg = read_config(conf_path)
                mounts = cfg.get("mounts") or default_mounts()
                metadata = collect_metadata()  # refresh host_ip/distro on reload
                _log("INFO", "config reloaded")
            except ConfigError as e:
                _log("ERROR", f"reload failed: {e}")
            _reload_requested = False

        sample = build_sample(mounts)
        buffered = buffer.drain()
        payload = build_payload(
            samples=buffered + [sample],
            hostname=hostname,
            metadata=metadata,
        )
        ok, status = post_payload(cfg["url"], cfg["token"], payload)

        if ok:
            if buffered:
                _log("INFO", f"flushed {len(buffered)} buffered samples")
            backoff = 0
            sleep_for = cfg["interval_seconds"]
        else:
            if status == 400:
                _log("ERROR", "server rejected payload (400) — dropping sample")
            elif status == 401:
                _log("ERROR", "token rejected (401) — check config + UI")
                buffer.push(sample)
            else:
                _log("WARN", f"POST failed (status={status}); buffering")
                for s in buffered:
                    buffer.push(s)
                buffer.push(sample)
            backoff = next_backoff(backoff)
            sleep_for = backoff

        slept = 0.0
        while slept < sleep_for and not _shutdown_requested and not _reload_requested:
            time.sleep(min(1.0, sleep_for - slept))
            slept += 1.0

    _log("INFO", "SIGTERM — flushing and exiting")
    final = buffer.drain()
    if final:
        post_payload(cfg["url"], cfg["token"],
                     build_payload(final, hostname, metadata))
    return 0


if __name__ == "__main__":
    conf = os.environ.get("STEWARD_AGENT_CONFIG", "/etc/steward-agent.conf")
    sys.exit(main_loop(conf))