#!/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.

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"
ACTIONS = ("update", "reboot", "drain", "undrain")
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:]}


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()}
    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())