FabledSteward/plugins/host_agent/routes.py

# plugins/host_agent/routes.py
from __future__ import annotations

import hashlib
from datetime import datetime, timezone
from typing import Any

from pathlib import Path

import secrets
from quart import Blueprint, current_app, jsonify, request, Response, render_template, redirect, url_for
from steward.auth.middleware import require_role
from steward.models.users import UserRole
from sqlalchemy import select, func
from datetime import timedelta

from steward.core.settings import public_base_url
from steward.models.hosts import Host
from steward.models.metrics import PluginMetric
from .models import HostAgentRegistration

host_agent_bp = Blueprint("host_agent", __name__, template_folder="templates")

SOURCE_MODULE = "host_agent"


def _hash_token(raw: str) -> str:
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()


def _parse_ts(ts: str) -> datetime:
    if ts.endswith("Z"):
        ts = ts[:-1] + "+00:00"
    return datetime.fromisoformat(ts)


def pick_host_address(current: str | None, reported: str | None) -> str | None:
    """Return the address to store on the Host, or None for no change.

    The agent-reported IP fills Host.address only when the current value is
    blank — an admin-typed address (DNS name, management IP) is never
    overwritten. The reported IP is always kept on the registration regardless,
    so admins can still see drift.
    """
    if reported and not (current or "").strip():
        return reported
    return None


def _expand_sample_to_metrics(
    sample: dict, host_name: str, recorded_at: datetime
) -> list[PluginMetric]:
    rows: list[PluginMetric] = []

    def row(metric: str, resource: str, value: float) -> None:
        rows.append(PluginMetric(
            source_module=SOURCE_MODULE,
            resource_name=resource,
            metric_name=metric,
            value=float(value),
            recorded_at=recorded_at,
        ))

    if sample.get("cpu_pct") is not None:
        row("cpu_pct", host_name, sample["cpu_pct"])

    mem = sample.get("mem") or {}
    if mem.get("total_bytes"):
        total = mem["total_bytes"]
        available = mem.get("available_bytes", 0)
        used_pct = 100.0 * (total - available) / total if total else 0.0
        row("mem_used_pct", host_name, used_pct)
        row("mem_available_bytes", host_name, available)
        row("swap_used_bytes", host_name, mem.get("swap_used_bytes", 0))

    load = sample.get("load") or {}
    for key in ("1m", "5m", "15m"):
        if key in load:
            row(f"load_{key}", host_name, load[key])

    if sample.get("uptime_secs") is not None:
        row("uptime_secs", host_name, sample["uptime_secs"])

    worst_pct = 0.0
    for disk in sample.get("storage") or []:
        total = disk.get("total_bytes", 0)
        used = disk.get("used_bytes", 0)
        pct = 100.0 * used / total if total else 0.0
        mount_res = f"{host_name}:{disk['mount']}"
        row("disk_used_pct", mount_res, pct)
        row("disk_used_bytes", mount_res, used)
        row("disk_total_bytes", mount_res, total)
        if pct > worst_pct:
            worst_pct = pct
    if sample.get("storage"):
        row("disk_used_pct_worst", host_name, worst_pct)

    # Per-core CPU. Sub-resources carry a ':' so the fleet widget filters them out.
    for i, core_pct in enumerate(sample.get("cpu_cores") or []):
        if core_pct is not None:
            row("cpu_pct", f"{host_name}:core{i}", core_pct)

    # Richer memory breakdown.
    if mem.get("cached_bytes") is not None:
        row("mem_cached_bytes", host_name, mem["cached_bytes"])
    if mem.get("buffers_bytes") is not None:
        row("mem_buffers_bytes", host_name, mem["buffers_bytes"])

    # Network throughput — per interface plus a host-level total for the fleet view.
    net_rx_total = net_tx_total = 0.0
    for iface in sample.get("net") or []:
        rx, tx = iface.get("rx_bps", 0.0), iface.get("tx_bps", 0.0)
        res = f"{host_name}:net:{iface['iface']}"
        row("net_rx_bps", res, rx)
        row("net_tx_bps", res, tx)
        net_rx_total += rx
        net_tx_total += tx
    if sample.get("net"):
        row("net_rx_bps", host_name, net_rx_total)
        row("net_tx_bps", host_name, net_tx_total)

    # Disk I/O — per device plus host-level total.
    dr_total = dw_total = 0.0
    for dev in sample.get("diskio") or []:
        rd, wr = dev.get("read_bps", 0.0), dev.get("write_bps", 0.0)
        res = f"{host_name}:diskio:{dev['device']}"
        row("disk_read_bps", res, rd)
        row("disk_write_bps", res, wr)
        dr_total += rd
        dw_total += wr
    if sample.get("diskio"):
        row("disk_read_bps", host_name, dr_total)
        row("disk_write_bps", host_name, dw_total)

    # Temperatures — per sensor plus the host max for at-a-glance.
    max_temp: float | None = None
    for t in sample.get("temps") or []:
        c = t.get("celsius")
        if c is None:
            continue
        row("temp_c", f"{host_name}:temp:{t['label']}", c)
        if max_temp is None or c > max_temp:
            max_temp = c
    if max_temp is not None:
        row("temp_c_max", host_name, max_temp)

    # Pressure stall information (PSI) — host-level gauges (mem/cpu/io some|full).
    for k, v in (sample.get("psi") or {}).items():
        if isinstance(v, (int, float)):
            row(f"psi_{k}", host_name, v)

    return rows


def _error(status: int, code: str, detail: str | None = None) -> tuple[Any, int]:
    body: dict = {"ok": False, "error": code}
    if detail:
        body["detail"] = detail
    return jsonify(body), status


@host_agent_bp.post("/ingest")
async def ingest():
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        return _error(401, "invalid_token")
    raw = auth[len("Bearer "):].strip()
    if not raw:
        return _error(401, "invalid_token")

    try:
        payload = await request.get_json(force=True)
    except Exception:
        return _error(400, "malformed_payload", "invalid JSON")
    if not isinstance(payload, dict) or "samples" not in payload:
        return _error(400, "malformed_payload", "missing 'samples'")

    samples = payload.get("samples") or []
    if not isinstance(samples, list) or not samples:
        return _error(400, "malformed_payload", "'samples' must be a non-empty list")

    token_hash = _hash_token(raw)
    async with current_app.db_sessionmaker() as session:
        reg = (await session.execute(
            select(HostAgentRegistration).where(
                HostAgentRegistration.token_hash == token_hash)
        )).scalar_one_or_none()
        if reg is None:
            return _error(401, "invalid_token")
        host = (await session.execute(
            select(Host).where(Host.id == reg.host_id)
        )).scalar_one_or_none()
        if host is None:
            return _error(401, "invalid_token")

        accepted = 0
        latest_ts: datetime | None = None
        for sample in samples:
            try:
                recorded_at = _parse_ts(sample["ts"])
            except (KeyError, ValueError):
                continue
            metrics = _expand_sample_to_metrics(sample, host.name, recorded_at)
            for m in metrics:
                session.add(m)
            accepted += 1
            if latest_ts is None or recorded_at > latest_ts:
                latest_ts = recorded_at

        if accepted == 0:
            await session.rollback()
            return _error(400, "malformed_payload", "no valid samples")

        md = payload.get("metadata") or {}
        changed = False
        if latest_ts and (reg.last_seen_at is None or latest_ts > reg.last_seen_at):
            reg.last_seen_at = latest_ts
            changed = True
        version = payload.get("agent_version")
        if version and reg.agent_version != version:
            reg.agent_version = version
            changed = True
        for field in ("kernel", "distro", "arch"):
            if field in md and getattr(reg, field) != md[field]:
                setattr(reg, field, md[field])
                changed = True
        host_ip = md.get("host_ip")
        if host_ip and reg.host_ip != host_ip:
            reg.host_ip = host_ip
            changed = True
        # Mirror the reported IP into Host.address only when it's blank.
        new_address = pick_host_address(host.address, host_ip)
        if new_address is not None:
            host.address = new_address
        if changed:
            reg.updated_at = datetime.now(timezone.utc)

        if latest_ts:
            skew = abs((datetime.now(timezone.utc) - latest_ts).total_seconds())
            if skew > 300:
                current_app.logger.warning(
                    "host_agent ingest: clock skew %.0fs for host=%s", skew, host.name)

        await session.commit()

    return jsonify({"ok": True, "samples_accepted": accepted}), 200


AGENT_SOURCE_PATH = Path(__file__).parent / "agent.py"


def _agent_version() -> str:
    try:
        for line in AGENT_SOURCE_PATH.read_text().splitlines():
            if line.startswith("AGENT_VERSION"):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    except OSError:
        pass
    return "unknown"


@host_agent_bp.get("/install.sh")
async def install_script():
    token = request.args.get("token", "")
    if not token:
        return _error(401, "invalid_token")
    async with current_app.db_sessionmaker() as session:
        reg = (await session.execute(
            select(HostAgentRegistration).where(
                HostAgentRegistration.token_hash == _hash_token(token))
        )).scalar_one_or_none()
        if reg is None:
            return _error(401, "invalid_token")
        host = (await session.execute(
            select(Host).where(Host.id == reg.host_id)
        )).scalar_one_or_none()
        if host is None:
            return _error(401, "invalid_token")

    url = public_base_url(request)

    rendered = await render_template(
        "install.sh.j2",
        url=url,
        token=token,
        agent_version=_agent_version(),
        host_name=host.name,
        host_address=host.address,
    )
    return Response(rendered, mimetype="text/plain")


@host_agent_bp.get("/agent.py")
async def agent_source():
    try:
        body = AGENT_SOURCE_PATH.read_text(encoding="utf-8")
    except OSError:
        return _error(500, "agent_missing")
    return Response(body, mimetype="text/x-python")


@host_agent_bp.get("/widget")
async def widget_table():
    """Fleet-glance table: one row per monitored host, latest metrics."""
    async with current_app.db_sessionmaker() as session:
        regs = (await session.execute(select(HostAgentRegistration))).scalars().all()
        host_ids = [r.host_id for r in regs]
        hosts = {
            h.id: h for h in (await session.execute(
                select(Host).where(Host.id.in_(host_ids))
            )).scalars().all()
        } if host_ids else {}

        subq = (
            select(
                PluginMetric.resource_name,
                PluginMetric.metric_name,
                func.max(PluginMetric.recorded_at).label("max_ts"),
            )
            .where(PluginMetric.source_module == SOURCE_MODULE)
            .group_by(PluginMetric.resource_name, PluginMetric.metric_name)
        ).subquery()
        latest_rows = (await session.execute(
            select(PluginMetric).join(
                subq,
                (PluginMetric.resource_name == subq.c.resource_name) &
                (PluginMetric.metric_name == subq.c.metric_name) &
                (PluginMetric.recorded_at == subq.c.max_ts),
            ).where(PluginMetric.source_module == SOURCE_MODULE)
        )).scalars().all()

    latest: dict[str, dict[str, float]] = {}
    for row in latest_rows:
        if ":" in row.resource_name:
            continue
        latest.setdefault(row.resource_name, {})[row.metric_name] = row.value

    rows = []
    for reg in regs:
        host = hosts.get(reg.host_id)
        if host is None:
            continue
        m = latest.get(host.name, {})
        rows.append({
            "host": host,
            "reg": reg,
            "cpu_pct": m.get("cpu_pct"),
            "mem_used_pct": m.get("mem_used_pct"),
            "disk_worst": m.get("disk_used_pct_worst"),
            "load_1m": m.get("load_1m"),
        })

    return await render_template("widget_table.html", rows=rows)


@host_agent_bp.get("/widget/history")
async def widget_history():
    host_id = request.args.get("host_id", "")
    hours = int(request.args.get("hours", "6"))
    cutoff = datetime.now(timezone.utc) - timedelta(hours=hours)

    async with current_app.db_sessionmaker() as session:
        host = (await session.execute(
            select(Host).where(Host.id == host_id))).scalar_one_or_none()
        if host is None:
            return _error(404, "not_found")
        points = (await session.execute(
            select(PluginMetric).where(
                PluginMetric.source_module == SOURCE_MODULE,
                PluginMetric.resource_name == host.name,
                PluginMetric.metric_name.in_(
                    ("cpu_pct", "mem_used_pct", "disk_used_pct_worst")),
                PluginMetric.recorded_at >= cutoff,
            ).order_by(PluginMetric.recorded_at)
        )).scalars().all()

    series: dict[str, list[dict]] = {"cpu_pct": [], "mem_used_pct": [], "disk_used_pct_worst": []}
    for p in points:
        series[p.metric_name].append({"t": p.recorded_at.isoformat(), "v": p.value})

    return await render_template("widget_history.html", host=host, series=series, hours=hours)


def _new_token_pair() -> tuple[str, str]:
    raw = secrets.token_urlsafe(32)
    return raw, _hash_token(raw)


@host_agent_bp.get("/settings/")
@require_role(UserRole.admin)
async def settings_list():
    async with current_app.db_sessionmaker() as session:
        regs = (await session.execute(select(HostAgentRegistration))).scalars().all()
        hosts_by_id = {
            h.id: h for h in (await session.execute(
                select(Host).where(Host.id.in_([r.host_id for r in regs])))).scalars().all()
        } if regs else {}

    new_token = request.args.get("new_token")
    new_host_id = request.args.get("host_id")
    install_url = None
    if new_token and new_host_id:
        install_url = f"{public_base_url(request)}/plugins/host_agent/install.sh?token={new_token}"

    return await render_template(
        "settings_list.html",
        registrations=[
            {"reg": r, "host": hosts_by_id.get(r.host_id)} for r in regs
        ],
        new_token=new_token,
        install_url=install_url,
    )


@host_agent_bp.post("/settings/add-host")
@require_role(UserRole.admin)
async def add_host():
    form = await request.form
    name = (form.get("name") or "").strip()
    address = (form.get("address") or "").strip()
    if not name:
        return _error(400, "missing_name")

    async with current_app.db_sessionmaker() as session:
        host = (await session.execute(
            select(Host).where(Host.name == name))).scalar_one_or_none()
        if host is None:
            host = Host(name=name, address=address)
            session.add(host)
            await session.flush()
        existing = (await session.execute(
            select(HostAgentRegistration).where(
                HostAgentRegistration.host_id == host.id))).scalar_one_or_none()
        if existing is not None:
            await session.rollback()
            return _error(400, "already_registered")

        raw, hashed = _new_token_pair()
        reg = HostAgentRegistration(host_id=host.id, token_hash=hashed)
        session.add(reg)
        await session.commit()
        host_id = host.id

    return redirect(url_for("host_agent.settings_list") + f"?new_token={raw}&host_id={host_id}")


@host_agent_bp.post("/settings/<host_id>/rotate-token")
@require_role(UserRole.admin)
async def rotate_token(host_id: str):
    async with current_app.db_sessionmaker() as session:
        reg = (await session.execute(
            select(HostAgentRegistration).where(
                HostAgentRegistration.host_id == host_id))).scalar_one_or_none()
        if reg is None:
            return _error(404, "not_found")
        raw, hashed = _new_token_pair()
        reg.token_hash = hashed
        reg.token_created_at = datetime.now(timezone.utc)
        await session.commit()
    return redirect(url_for("host_agent.settings_list") + f"?new_token={raw}&host_id={host_id}")


@host_agent_bp.post("/settings/<host_id>/delete")
@require_role(UserRole.admin)
async def delete_registration(host_id: str):
    async with current_app.db_sessionmaker() as session:
        reg = (await session.execute(
            select(HostAgentRegistration).where(
                HostAgentRegistration.host_id == host_id))).scalar_one_or_none()
        if reg is not None:
            await session.delete(reg)
            await session.commit()
    return redirect(url_for("host_agent.settings_list"))