#!/usr/bin/env python3
"""Steward maintenance helper (#5766): runs as root, started by steward-maint.path
when the host agent leaves a request, and does only what ACTIONS lists.

The agent (an unprivileged user under NoNewPrivileges) can't update or reboot
its host itself. It writes `<id>.json` = {"id", "action", ["node"]} into
REQUESTS, the one directory it may write. This helper reads each request
without following links, refuses anything outside the list, deletes it, does
the action, and writes the outcome to RESULTS (root's own directory, which the
agent reads back and reports to Steward). A stolen agent token can therefore
make its host update, reboot or drain, and nothing else.

One action only reads: `drives` (#5789) asks smartctl for each disk's health,
which needs root, and writes a summary to DRIVES, never a result: it is not a
maintenance job.

Stdlib only; served by Steward at /plugins/host_agent/files/maint.py.
"""
from __future__ import annotations

import json
import os
import re
import shutil
import subprocess
import sys
import time
from datetime import datetime, timezone

REQUESTS = "/var/lib/steward-agent/requests"
RESULTS = "/var/lib/steward-maint/results"
DRIVES = "/var/lib/steward-maint/drives.json"
ACTIONS = ("update", "reboot", "drain", "undrain", "drives")
REQUEST_MAX_BYTES = 4096
RESULTS_KEPT_SECONDS = 7 * 86400
UPDATE_DEADLINE_SECONDS = 3000
COMMAND_DEADLINE_SECONDS = 600
OUTPUT_KEPT_CHARS = 2000
_ID = re.compile(r"^[A-Za-z0-9-]{1,64}$")
_NODE = re.compile(r"^[a-z0-9]{25}$")


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def read_request(path: str):
    """The request at `path` as a dict, or None when it is not a small regular
    file holding a valid request. Never follows a link the agent planted."""
    try:
        fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
    except OSError:
        return None
    try:
        st = os.fstat(fd)
        if not (st.st_mode & 0o170000 == 0o100000) or st.st_size > REQUEST_MAX_BYTES:
            return None
        raw = os.read(fd, REQUEST_MAX_BYTES + 1)
    finally:
        os.close(fd)
    try:
        req = json.loads(raw.decode("utf-8"))
    except (UnicodeDecodeError, ValueError):
        return None
    return req if valid(req) else None


def valid(req) -> bool:
    """A request names a known action, has a plain id, and a node id only in
    the swarm's own form."""
    if not isinstance(req, dict) or req.get("action") not in ACTIONS:
        return False
    if not isinstance(req.get("id"), str) or not _ID.match(req["id"]):
        return False
    node = req.get("node")
    return node is None or (isinstance(node, str) and bool(_NODE.match(node)))


def run(cmd: list, deadline: int = COMMAND_DEADLINE_SECONDS, env: dict | None = None):
    """(ok, output tail) of a command, bounded by `deadline`."""
    try:
        done = subprocess.run(cmd, capture_output=True, text=True, timeout=deadline,
                              env=dict(os.environ, LANG="C", LC_ALL="C", **(env or {})))
    except subprocess.TimeoutExpired:
        return False, f"{cmd[0]} didn't finish in {deadline} s"
    except OSError as exc:
        return False, str(exc)
    out = (done.stdout or "") + (done.stderr or "")
    return done.returncode == 0, out[-OUTPUT_KEPT_CHARS:]


def reboot_required() -> bool | None:
    if os.path.exists("/var/run/reboot-required"):
        return True
    if shutil.which("apt-get"):
        return False
    if shutil.which("needs-restarting"):
        try:
            return subprocess.run(["needs-restarting", "-r"], capture_output=True,
                                  timeout=60).returncode == 1
        except (OSError, subprocess.TimeoutExpired):
            return None
    return None


def upgraded_count(output: str) -> int | None:
    """Packages an apt or dnf run changed, from its summary line."""
    m = re.search(r"(\d+) upgraded, (\d+) newly installed", output)       # apt
    if m:
        return int(m.group(1)) + int(m.group(2))
    counts = re.findall(r"^\s*(?:Upgrade|Install)\s+(\d+) Packages?", output, re.MULTILINE)  # dnf
    return sum(int(c) for c in counts) if counts else None


def do_update() -> dict:
    """Upgrade every package, keeping existing config files and leaving
    service restarts to a reboot: needrestart only lists, so the Docker daemon
    isn't restarted under its containers mid-upgrade."""
    if shutil.which("apt-get"):
        env = {"DEBIAN_FRONTEND": "noninteractive", "NEEDRESTART_MODE": "l"}
        ok, out = run(["apt-get", "update"], env=env)
        if not ok:
            return {"ok": False, "detail": "apt-get update failed: " + out}
        ok, out = run(["apt-get", "-y", "-o", "Dpkg::Options::=--force-confdef",
                       "-o", "Dpkg::Options::=--force-confold", "dist-upgrade"],
                      deadline=UPDATE_DEADLINE_SECONDS, env=env)
        if not ok:
            return {"ok": False, "detail": "apt-get dist-upgrade failed: " + out}
        changed = upgraded_count(out)
        run(["apt-get", "-y", "autoremove"], env=env)
    elif shutil.which("dnf"):
        ok, out = run(["dnf", "-y", "upgrade", "--refresh"], deadline=UPDATE_DEADLINE_SECONDS)
        if not ok:
            return {"ok": False, "detail": "dnf upgrade failed: " + out}
        changed = upgraded_count(out)
    else:
        return {"ok": False, "detail": "no apt-get or dnf on this host"}
    return {"ok": True, "packages": changed, "reboot_required": reboot_required(),
            "detail": out[-500:]}


def own_node() -> str | None:
    ok, out = run(["docker", "info", "--format", "{{.Swarm.NodeID}}"], deadline=30)
    node = out.strip() if ok else ""
    return node if _NODE.match(node) else None


def do_node(availability: str, node: str | None) -> dict:
    """Set a swarm node's availability (drain before a reboot, active after).
    Only a manager can; a worker's helper says so and Steward asks a manager."""
    if not shutil.which("docker"):
        return {"ok": False, "detail": "no docker on this host"}
    node = node or own_node()
    if node is None:
        return {"ok": False, "detail": "this host is not in a swarm"}
    ok, out = run(["docker", "node", "update", "--availability", availability, node], deadline=60)
    return {"ok": ok, "node": node, "detail": out.strip()[-500:]}


# ── drive health (#5789) ─────────────────────────────────────────────────────
# smartctl ≥ 7.0 answers in JSON, the same for SATA, SAS and NVMe drives on any
# distribution; an older one is reported, never parsed as text. The helper
# finds the drives itself: the request names none.

SMARTCTL_MIN = (7, 0)
SMARTCTL_DEADLINE_SECONDS = 60
DRIVES_MAX = 32
# smartctl's exit status is a bit mask: 1 the command line, 2 the device
# couldn't be opened. Higher bits describe the disk, and its JSON still holds.
_SMARTCTL_UNREAD = 0b11
# ATA attributes counted as bad sectors.
_ATA_BAD = {5: "reallocated", 197: "pending", 198: "offline_uncorrectable"}


def smartctl_version():
    """(major, minor) of the installed smartctl, or None when there is none."""
    if not shutil.which("smartctl"):
        return None
    ok, out = run(["smartctl", "--version"], deadline=30)
    m = re.search(r"smartctl (\d+)\.(\d+)", out)
    return (int(m.group(1)), int(m.group(2))) if m else (0, 0)


def smartctl_json(args: list):
    """(exit status, parsed JSON or None) of one smartctl call."""
    try:
        done = subprocess.run(["smartctl", "--json", *args], capture_output=True, text=True,
                              timeout=SMARTCTL_DEADLINE_SECONDS,
                              env=dict(os.environ, LANG="C", LC_ALL="C"))
    except (subprocess.TimeoutExpired, OSError):
        return None, None
    try:
        return done.returncode, json.loads(done.stdout)
    except ValueError:
        return done.returncode, None


def _int(v):
    return v if isinstance(v, int) and not isinstance(v, bool) else None


def drive_name(name: str, kind) -> str:
    """A drive's name, unique on its host. Disks behind a RAID controller
    share one path (/dev/bus/0) and differ only by the slot in their type
    (megaraid,3; cciss,0), so the slot is part of the name (#5806)."""
    base = os.path.basename(name)
    if name.startswith("/dev/bus/"):
        base = "bus" + base
    if isinstance(kind, str) and "," in kind:
        base += "-" + re.sub(r"[^A-Za-z0-9]", "", kind.split("+")[-1])
    return base


def drive_summary(name: str, status, data, kind=None) -> dict:
    """The few figures Steward reads from one drive's `smartctl -a` JSON."""
    out = {"device": drive_name(name, kind)}
    if not isinstance(data, dict) or status is None or status & _SMARTCTL_UNREAD:
        messages = ((data.get("smartctl") or {}).get("messages") or []) if isinstance(data, dict) else []
        text = "; ".join(str(m.get("string", "")) for m in messages if isinstance(m, dict))
        out["error"] = (text or "smartctl couldn't read it")[:200]
        return out
    dev = data.get("device") or {}
    out["protocol"] = str(dev.get("protocol") or "")[:16]
    out["model"] = str(data.get("model_name") or data.get("scsi_model_name") or "")[:80]
    cap = (data.get("user_capacity") or {}).get("bytes")
    if _int(cap):
        out["capacity_bytes"] = cap
    support = data.get("smart_support") or {}
    if support.get("available") is False or support.get("enabled") is False:
        out["available"] = False             # a virtual disk, or SMART switched off
        return out
    passed = (data.get("smart_status") or {}).get("passed")
    if isinstance(passed, bool):
        out["passed"] = passed
    temp = _int((data.get("temperature") or {}).get("current"))
    if temp is not None:
        out["temp_c"] = temp
    hours = _int((data.get("power_on_time") or {}).get("hours"))
    if hours is not None:
        out["power_on_hours"] = hours
    nvme = data.get("nvme_smart_health_information_log") or {}
    for src, dst in (("critical_warning", "critical_warning"), ("percentage_used", "life_used_pct"),
                     ("media_errors", "media_errors"), ("available_spare", "spare_pct"),
                     ("available_spare_threshold", "spare_threshold_pct")):
        if _int(nvme.get(src)) is not None:
            out[dst] = nvme[src]
    for attr in (data.get("ata_smart_attributes") or {}).get("table") or []:
        key = _ATA_BAD.get(attr.get("id")) if isinstance(attr, dict) else None
        raw = _int((attr.get("raw") or {}).get("value")) if key else None
        if raw is not None:
            out[key] = raw
    used = _int((data.get("endurance_used") or {}).get("current_percent"))          # ATA SSDs
    used = used if used is not None else _int(data.get("scsi_percentage_used_endurance_indicator"))
    if used is not None and "life_used_pct" not in out:
        out["life_used_pct"] = used
    if _int(data.get("scsi_grown_defect_list")) is not None:
        out["grown_defects"] = data["scsi_grown_defect_list"]
    return out


def do_drives() -> dict:
    """Each drive's health, as {"checked_at", "smartctl", "drives"} or
    {"checked_at", "error"} when smartctl is missing or too old."""
    base = {"checked_at": now_iso()}
    version = smartctl_version()
    if version is None:
        return {**base, "error": "smartctl not installed"}
    base["smartctl"] = "%d.%d" % version
    if version < SMARTCTL_MIN:
        return {**base, "error": "smartctl %d.%d is too old: drive health needs 7.0 or newer" % version}
    _, scan = smartctl_json(["--scan-open"])
    devices = (scan or {}).get("devices") if isinstance(scan, dict) else None
    if not isinstance(devices, list):
        return {**base, "error": "smartctl couldn't list the drives"}
    drives = []
    for d in devices[:DRIVES_MAX]:
        name, kind = (d.get("name"), d.get("type")) if isinstance(d, dict) else (None, None)
        if not isinstance(name, str) or not name.startswith("/dev/"):
            continue
        args = ["-a", name] if not isinstance(kind, str) else ["-a", "-d", kind, name]
        status, data = smartctl_json(args)
        drives.append(drive_summary(name, status, data, kind))
    return {**base, "drives": drives}


def write_drives(report: dict, path: str | None = None) -> None:
    """DRIVES, written whole and readable by the agent."""
    path = path or DRIVES
    tmp = path + ".tmp"
    with open(tmp, "w") as fh:
        json.dump(report, fh)
    os.chmod(tmp, 0o644)
    os.replace(tmp, path)


def write_result(result: dict) -> None:
    """RESULTS/<id>.json, written whole (a temp file renamed into place)."""
    os.makedirs(RESULTS, mode=0o755, exist_ok=True)
    final = os.path.join(RESULTS, result["id"] + ".json")
    tmp = final + ".tmp"
    with open(tmp, "w") as fh:
        json.dump(result, fh)
    os.chmod(tmp, 0o644)
    os.replace(tmp, final)


def already_handled(job_id: str) -> bool:
    """A request the agent sent again, for a job already run, is not run twice."""
    return os.path.exists(os.path.join(RESULTS, job_id + ".json"))


def handle(req: dict) -> dict:
    base = {"id": req["id"], "action": req["action"], "started_at": now_iso()}
    if req["action"] == "drives":               # a read, not a job: no result
        write_drives(do_drives())
        return {**base, "state": "done"}
    write_result({**base, "state": "running"})
    action = req["action"]
    if action == "update":
        outcome = do_update()
    elif action in ("drain", "undrain"):
        outcome = do_node("drain" if action == "drain" else "active", req.get("node"))
    else:                                            # reboot
        result = {**base, "state": "rebooting", "ok": True, "finished_at": now_iso()}
        write_result(result)
        os.sync()
        run(["systemctl", "reboot"], deadline=60)
        return result
    result = {**base, **outcome, "state": "done" if outcome.get("ok") else "failed",
              "finished_at": now_iso()}
    write_result(result)
    return result


def prune_results(now: float | None = None) -> None:
    now = time.time() if now is None else now
    try:
        names = os.listdir(RESULTS)
    except OSError:
        return
    for name in names:
        path = os.path.join(RESULTS, name)
        try:
            if now - os.lstat(path).st_mtime > RESULTS_KEPT_SECONDS:
                os.unlink(path)
        except OSError:
            pass


def pending_requests() -> list:
    """Request files, oldest first."""
    try:
        names = [n for n in os.listdir(REQUESTS) if n.endswith(".json")]
    except OSError:
        return []
    paths = [os.path.join(REQUESTS, n) for n in names]

    def age(p):
        try:
            return os.lstat(p).st_mtime
        except OSError:
            return 0.0
    return sorted(paths, key=age)


def main() -> int:
    prune_results()
    for path in pending_requests():
        req = read_request(path)
        try:
            os.unlink(path)                 # whatever it held: never re-read in a loop
        except OSError:
            pass
        if req is None:
            sys.stderr.write(f"[steward-maint] ignored {os.path.basename(path)}: not a valid request\n")
            continue
        if already_handled(req["id"]):
            continue
        result = handle(req)
        sys.stderr.write(f"[steward-maint] {req['action']} {req['id']}: {result['state']}\n")
        if result["state"] == "rebooting":
            break                           # the rest waits for the next boot
    return 0


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