dashboard/app/utils/systemd.py

import json
import re
import time
from typing import Any

from app.utils.subprocess import run_async

UNIT_RE = re.compile(r"^[A-Za-z0-9@:_.\-+]+\.(service|socket|timer|target|path|slice)$")
ACTIONS = ("start", "stop", "restart", "enable", "disable")

_enabled_cache: dict[str, str] | None = None
_enabled_cache_at = 0.0
_ENABLED_TTL = 30.0

_DETAIL_PROPS = (
    "ActiveState,SubState,LoadState,UnitFileState,Description,MainPID,"
    "ExecMainStartTimestamp,NRestarts,FragmentPath,Result"
)


async def _systemctl(*args: str, privileged: bool = False) -> str:
    """Run a systemctl command and return its stdout.

    Args:
        *args: systemctl subcommand and options, e.g. ("show", "foo.service").
        privileged: run via sudo. Set for verbs that modify state (start,
            stop, enable, ...); plain attempts just get rejected by
            systemd and spam the journal with auth failures.

    Returns:
        The decoded stdout.

    Raises:
        RuntimeError: if systemctl exits non-zero (or cannot be spawned);
            the message is its stderr (or "systemctl <args> failed" when
            stderr is empty).
    """
    cmd = (["sudo", "systemctl", *args] if privileged else ["systemctl", *args])
    rc, out, err = await run_async(cmd)
    if rc == 0:
        return out
    raise RuntimeError(err.strip() or f"systemctl {' '.join(args)} failed")


async def _enabled_map(force: bool = False) -> dict[str, str]:
    """Map unit name to enabled-state (enabled, disabled, static, ...).

    The result of `systemctl list-unit-files --type=service --output=json`
    is cached module-wide for 30 s so fast polls don't re-run it;
    unit_action() invalidates the cache after enable/disable.

    Args:
        force: bypass the cache and re-query.

    Returns:
        A unit-name to state-string mapping (may include units that are
        not currently active).
    """
    global _enabled_cache, _enabled_cache_at
    now = time.monotonic()
    if not force and _enabled_cache is not None and now - _enabled_cache_at < _ENABLED_TTL:
        return _enabled_cache
    out = await _systemctl("list-unit-files", "--type=service", "--output=json")
    rows: list[dict[str, Any]] = json.loads(out)
    m: dict[str, str] = {e["unit_file"]: e["state"] for e in rows}
    _enabled_cache = m
    _enabled_cache_at = now
    return m


async def unit_list() -> list[dict[str, str]]:
    """List all service units with their runtime and enabled state.

    Merges `systemctl list-units --all --output=json` (currently known
    units) with the enabled-state map, so units that are configured but
    not active still appear (with placeholder load/active/sub values).

    Returns:
        One row per unit (name, load, active, sub, desc, enabled),
        sorted by unit name.
    """
    out = await _systemctl("list-units", "--type=service", "--all", "--output=json")
    enabled = await _enabled_map()
    units: dict[str, dict[str, str]] = {}
    for e in json.loads(out):
        units[e["unit"]] = {
            "name": e["unit"],
            "load": e["load"],
            "active": e["active"],
            "sub": e["sub"],
            "desc": e.get("description", ""),
            "enabled": enabled.get(e["unit"], ""),
        }
    for name, state in enabled.items():
        if name not in units:
            units[name] = {
                "name": name,
                "load": "",
                "active": "inactive",
                "sub": "dead",
                "desc": "",
                "enabled": state,
            }
    return sorted(units.values(), key=lambda u: u["name"])


async def unit_detail(name: str) -> dict[str, str]:
    """Fetch the detail properties of one unit via `systemctl show`.

    Args:
        name: unit name, must match UNIT_RE.

    Returns:
        The requested properties (see _DETAIL_PROPS) as a key to value
        mapping, empty values included.

    Raises:
        ValueError: if the unit name is not a valid systemd unit name.
        RuntimeError: if systemctl fails.
    """
    if not UNIT_RE.match(name):
        raise ValueError("invalid unit name")
    out = await _systemctl("show", name, f"-p{_DETAIL_PROPS}")
    props: dict[str, str] = {}
    for line in out.splitlines():
        if "=" in line:
            k, _, v = line.partition("=")
            props[k] = v
    return props


async def unit_action(name: str, action: str) -> str:
    """Perform a state-changing verb on a unit (via sudo).

    enable/disable also invalidate the module-level enabled-state cache
    so the next unit_list() reflects the change immediately.

    Args:
        name: unit name, must match UNIT_RE.
        action: one of ACTIONS (start, stop, restart, enable, disable).

    Returns:
        The (usually empty) stdout of the systemctl call.

    Raises:
        ValueError: if the unit name or action is invalid.
        RuntimeError: if systemctl fails (e.g. unit does not exist).
    """
    if not UNIT_RE.match(name):
        raise ValueError("invalid unit name")
    if action not in ACTIONS:
        raise ValueError("invalid action")
    out = await _systemctl(action, name, privileged=True)
    if action in ("enable", "disable"):
        global _enabled_cache, _enabled_cache_at
        _enabled_cache = None
        _enabled_cache_at = 0.0
    return out


async def system_state() -> str:
    """Overall systemd state (running, degraded, ..., or "unknown").

    Returns:
        The `systemctl is-system-running` state, or "unknown" when the
        call fails (e.g. inside a container).
    """
    try:
        return (await _systemctl("is-system-running")).strip() or "unknown"
    except RuntimeError:
        return "unknown"