#!/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_codes: tuple = (0,)):
    """(ok, output tail) of a command, bounded by `deadline`. `ok_codes` are
    the exit codes that mean it worked (zypper says more than 0)."""
    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 in ok_codes, out[-OUTPUT_KEPT_CHARS:]


REBOOT_MARKER = "/var/run/reboot-required"                 # Debian and Ubuntu
KERNEL_MODULES = "/usr/lib/modules"
# zypper's own exit codes for "done, and": 102 a reboot is needed, 103 zypper
# updated itself and would run again (its next run picks up the rest).
ZYPPER_OK = (0, 102, 103)


def reboot_required() -> bool | None:
    """Whether the host needs a reboot after an update, asked the way its
    package manager says it; None when it can't say."""
    if os.path.exists(REBOOT_MARKER):
        return True
    if shutil.which("apt-get"):
        return False
    try:
        if shutil.which("needs-restarting"):                   # dnf, yum
            return subprocess.run(["needs-restarting", "-r"], capture_output=True,
                                  timeout=60).returncode == 1
        if shutil.which("zypper"):
            code = subprocess.run(["zypper", "needs-rebooting"], capture_output=True,
                                  timeout=60).returncode
            return {0: False, 102: True}.get(code)
    except (OSError, subprocess.TimeoutExpired):
        return None
    if shutil.which("pacman"):
        # An upgraded kernel removes the running one's modules.
        return not os.path.isdir(os.path.join(KERNEL_MODULES, os.uname().release))
    return None


def upgraded_count(output: str) -> int | None:
    """Packages an update changed, from the manager's summary line."""
    m = re.search(r"(\d+) upgraded, (\d+) newly installed", output)       # apt
    if m:
        return int(m.group(1)) + int(m.group(2))
    m = re.search(r"(\d+) packages? to upgrade(?:, (\d+) new)?", output)  # zypper
    if m:
        return int(m.group(1)) + int(m.group(2) or 0)
    m = re.search(r"^Packages \((\d+)\)", output, re.MULTILINE)           # pacman
    if m:
        return int(m.group(1))
    counts = re.findall(r"^\s*(?:Upgrade|Install)\s+(\d+) Packages?", output, re.MULTILINE)  # dnf, yum
    return sum(int(c) for c in counts) if counts else None


# The same as the agent's: this helper runs as root on its own and imports
# nothing of the agent's, so it keeps its own copy.
def _os_release_id(path: str = "/etc/os-release") -> str:
    try:
        with open(path) as fh:
            for line in fh:
                if line.startswith("ID="):
                    return line.split("=", 1)[1].strip().strip('"')
    except OSError:
        pass
    return ""


def update_steps(manager: str) -> list:
    """The commands an update runs on a host with `manager`, in order: each
    (what it is, command, deadline, env, ok exit codes). The last one is the upgrade
    whose output is counted. Existing config files are kept, and service
    restarts are left to a reboot: needrestart only lists, so the Docker
    daemon isn't restarted under its containers mid-upgrade."""
    long = UPDATE_DEADLINE_SECONDS
    if manager == "apt":
        env = {"DEBIAN_FRONTEND": "noninteractive", "NEEDRESTART_MODE": "l"}
        return [("apt-get update", ["apt-get", "update"], COMMAND_DEADLINE_SECONDS, env, (0,)),
                ("apt-get dist-upgrade", ["apt-get", "-y", "-o", "Dpkg::Options::=--force-confdef",
                                          "-o", "Dpkg::Options::=--force-confold", "dist-upgrade"],
                 long, env, (0,))]
    if manager == "dnf":
        return [("dnf upgrade", ["dnf", "-y", "upgrade", "--refresh"], long, None, (0,))]
    if manager == "yum":
        return [("yum update", ["yum", "-y", "update"], long, None, (0,))]
    if manager == "zypper":
        # Tumbleweed is a rolling release: it upgrades with dup, not update.
        upgrade = "dup" if _os_release_id() == "opensuse-tumbleweed" else "update"
        return [("zypper refresh", ["zypper", "--non-interactive", "refresh"],
                 COMMAND_DEADLINE_SECONDS, None, ZYPPER_OK),
                ("zypper " + upgrade, ["zypper", "--non-interactive", upgrade, "--auto-agree-with-licenses"],
                 long, None, ZYPPER_OK)]
    if manager == "pacman":
        # Packages from the AUR aren't updated: no AUR helper runs as root.
        return [("pacman -Syu", ["pacman", "-Syu", "--noconfirm"], long, None, (0,))]
    return []


# Which manager a host has; dnf before yum, since dnf hosts often keep a yum
# that is dnf underneath.
UPDATE_MANAGERS = (("apt", "apt-get"), ("dnf", "dnf"), ("yum", "yum"), ("zypper", "zypper"),
                   ("pacman", "pacman"))


def host_manager() -> str | None:
    for manager, binary in UPDATE_MANAGERS:
        if shutil.which(binary):
            return manager
    return None


def do_update() -> dict:
    """Upgrade every package with whichever package manager the host has."""
    manager = host_manager()
    if manager is None:
        return {"ok": False, "detail": "no package manager this helper knows "
                                       "(apt, dnf, yum, zypper, pacman) on this host"}
    out = ""
    for label, cmd, deadline, env, ok_codes in update_steps(manager):
        ok, out = run(cmd, deadline=deadline, env=env, ok_codes=ok_codes)
        if not ok:
            return {"ok": False, "detail": f"{label} failed: " + out}
    changed = upgraded_count(out)
    if manager == "apt":
        run(["apt-get", "-y", "autoremove"], env={"DEBIAN_FRONTEND": "noninteractive"})
    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())