rigdoctor/src/rigdoctor/core/fixes.py

"""Apply runtime-reversible system tunables (M6) — a limited, consent-gated exception to
the read-only stance (D9, amended by D22).

Only safe settings that take effect immediately, need no reboot, and revert on reboot are
applyable here: CPU governor, NVIDIA persistence mode, PCIe ASPM policy, vm.swappiness, and
Transparent HugePages. Each is set by a single privileged command (one pkexec prompt). The
chosen value is validated against the live options before building the command, and writes go
to sysfs / procfs (or `nvidia-smi`) — never the GRUB cmdline or a persistent config file.
Riskier fixes (GRUB-based PCIe ASPM-off, CPU mitigations) stay suggestion-only.
"""

from __future__ import annotations

import os
import shlex
import shutil
import subprocess
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path


@dataclass
class Tunable:
    id: str
    label: str            # e.g. "CPU governor"
    options: list[str]    # selectable values (live, from the system)
    current: str | None   # the value in effect now (preselect this in the dropdown)
    note: str = ""        # caveat shown by the control, e.g. "resets on reboot"


def _read(path: str) -> str | None:
    try:
        return Path(path).read_text()
    except OSError:
        return None


def _bracketed(text: str) -> tuple[list[str], str | None]:
    """Parse a sysfs 'a [b] c' enum into (options, active)."""
    options = [tok.strip("[]") for tok in text.split()]
    active = next((tok.strip("[]") for tok in text.split() if tok.startswith("[")), None)
    return options, active


# --- individual tunables: a state reader + a command builder per id -------------------

_GOV = "/sys/devices/system/cpu"


def _cpu_governor() -> Tunable | None:
    cur = _read(f"{_GOV}/cpu0/cpufreq/scaling_governor")
    if cur is None:
        return None
    avail = _read(f"{_GOV}/cpu0/cpufreq/scaling_available_governors")
    options = avail.split() if avail and avail.strip() else ["performance", "powersave", "schedutil"]
    return Tunable("cpu_governor", "CPU governor", options, cur.strip(), "applies now; resets on reboot")


def _cpu_governor_cmd(value: str) -> list[str]:
    return ["/bin/sh", "-c",
            f'for f in {_GOV}/cpu*/cpufreq/scaling_governor; do echo {shlex.quote(value)} > "$f"; done']


def _nvidia_persistence() -> Tunable | None:
    if shutil.which("nvidia-smi") is None:
        return None
    try:
        proc = subprocess.run(
            ["nvidia-smi", "--query-gpu=persistence_mode", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10,
        )
    except (subprocess.SubprocessError, OSError):
        return None
    state = proc.stdout.strip().splitlines()[0].strip().lower() if proc.stdout.strip() else ""
    current = "Enabled" if state.startswith("enabled") else ("Disabled" if state.startswith("disabled") else None)
    return Tunable("nvidia_persistence", "NVIDIA persistence mode", ["Enabled", "Disabled"], current,
                   "resets on reboot (enable nvidia-persistenced to persist)")


def _nvidia_persistence_cmd(value: str) -> list[str]:
    return ["nvidia-smi", "-pm", "1" if value == "Enabled" else "0"]


def _pcie_aspm() -> Tunable | None:
    text = _read("/sys/module/pcie_aspm/parameters/policy")
    if not text:
        return None
    options, active = _bracketed(text)
    return Tunable("pcie_aspm", "PCIe ASPM policy", options, active, "applies now; resets on reboot")


def _pcie_aspm_cmd(value: str) -> list[str]:
    return ["/bin/sh", "-c", f'echo {shlex.quote(value)} > /sys/module/pcie_aspm/parameters/policy']


def _swappiness() -> Tunable | None:
    text = _read("/proc/sys/vm/swappiness")
    if text is None or not text.strip().isdigit():
        return None
    cur = text.strip()
    options = ["0", "10", "30", "60", "100"]
    if cur not in options:
        options = sorted(set(options) | {cur}, key=int)
    return Tunable("swappiness", "vm.swappiness", options, cur, "applies now; resets on reboot")


def _swappiness_cmd(value: str) -> list[str]:
    return ["/bin/sh", "-c", f'echo {shlex.quote(value)} > /proc/sys/vm/swappiness']


def _thp() -> Tunable | None:
    text = _read("/sys/kernel/mm/transparent_hugepage/enabled")
    if not text:
        return None
    options, active = _bracketed(text)
    return Tunable("thp", "Transparent HugePages", options, active, "applies now; resets on reboot")


def _thp_cmd(value: str) -> list[str]:
    return ["/bin/sh", "-c", f'echo {shlex.quote(value)} > /sys/kernel/mm/transparent_hugepage/enabled']


_TUNABLES: dict[str, tuple[Callable[[], Tunable | None], Callable[[str], list[str]]]] = {
    "cpu_governor": (_cpu_governor, _cpu_governor_cmd),
    "nvidia_persistence": (_nvidia_persistence, _nvidia_persistence_cmd),
    "pcie_aspm": (_pcie_aspm, _pcie_aspm_cmd),
    "swappiness": (_swappiness, _swappiness_cmd),
    "thp": (_thp, _thp_cmd),
}


# --- public API -----------------------------------------------------------------------

def get_tunable(fix_id: str) -> Tunable | None:
    """Live state (options + current value) for a fix id, or None if not applicable here."""
    fns = _TUNABLES.get(fix_id)
    return fns[0]() if fns else None


def apply_command(fix_id: str, value: str) -> list[str] | None:
    """The privileged command to set fix_id=value, or None if unknown/invalid.

    The value is validated against the *live* options, so only a real, currently-available
    setting can ever be turned into a command.
    """
    fns = _TUNABLES.get(fix_id)
    if not fns:
        return None
    state = fns[0]()
    if state is None or value not in state.options:
        return None
    return fns[1](value)


def _elevate(cmd: list[str]) -> list[str]:
    prog = shutil.which(cmd[0]) or cmd[0]  # pkexec needs an absolute program path
    cmd = [prog, *cmd[1:]]
    if os.geteuid() == 0:
        return cmd
    if shutil.which("pkexec"):
        return ["pkexec", *cmd]
    if shutil.which("sudo"):
        return ["sudo", *cmd]
    return cmd  # no escalation available — will likely fail, surfaced to the caller


def apply(fix_id: str, value: str) -> tuple[int, str]:
    """Apply fix_id=value via a single elevated command. Returns (exit_code, output)."""
    cmd = apply_command(fix_id, value)
    if cmd is None:
        return (1, f"Unknown or unavailable setting: {fix_id}={value}")
    try:
        proc = subprocess.run(_elevate(cmd), capture_output=True, text=True, timeout=120)
        return (proc.returncode, proc.stdout + proc.stderr)
    except (subprocess.SubprocessError, OSError) as exc:
        return (1, str(exc))