dashboard/app/systemd/units.py

import asyncio
import re
import time

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 _run(cmd: list[str]) -> tuple[int, str, str]:
    """Run a command, capturing stdout and stderr.

    Args:
        cmd: program and arguments.

    Returns:
        (returncode, stdout, stderr), all decoded; a missing returncode
        (should not happen) is reported as 0.
    """
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    out, err = await proc.communicate()
    return proc.returncode or 0, out.decode(errors="replace"), err.decode(errors="replace")


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; 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(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` 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
    files = await _systemctl("list-unit-files", "--type=service", "--no-legend", "--plain")
    m: dict[str, str] = {}
    for line in files.splitlines():
        parts = line.split(None, 2)
        if len(parts) < 2:
            continue
        m[parts[0]] = parts[1].strip()
    _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` (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", "--no-legend", "--plain"
    )
    enabled = await _enabled_map()
    units: dict[str, dict[str, str]] = {}
    for line in out.splitlines():
        parts = line.split(None, 4)
        if len(parts) < 4:
            continue
        name, load, active, sub = parts[0], parts[1], parts[2], parts[3]
        desc = parts[4] if len(parts) > 4 else ""
        units[name] = {
            "name": name,
            "load": load,
            "active": active,
            "sub": sub,
            "desc": desc,
            "enabled": enabled.get(name, ""),
        }
    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"