fan-aggressor-helper

#!/usr/bin/env python3

from __future__ import annotations

import sys
import os
import json
from pathlib import Path

CONFIG_FILE = Path("/etc/fan-aggressor/config.json")

ALLOWED_ACTIONS = {"save-config", "restart-service", "apply-cpu-power"}

ALLOWED_CONFIG_KEYS = {
    "cpu_fan_offset", "gpu_fan_offset", "enabled", "poll_interval",
    "hybrid_mode", "temp_threshold_engage", "temp_threshold_disengage",
    "cpu_governor", "cpu_turbo_enabled", "cpu_epp", "cpu_platform_profile",
    "link_offsets", "nekroctl_path", "failsafe_mode",
    "cpu_rapl_pl1_w", "cpu_rapl_pl2_w", "cpu_max_freq_mhz"
}


def save_config():
    content = sys.stdin.read()
    config = json.loads(content)
    sanitized = {k: v for k, v in config.items() if k in ALLOWED_CONFIG_KEYS}
    CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = CONFIG_FILE.with_suffix(".tmp")
    fd = os.open(str(tmp_path), os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644)
    with os.fdopen(fd, "w") as f:
        json.dump(sanitized, f, indent=2)
    os.replace(tmp_path, CONFIG_FILE)


def restart_service():
    import subprocess
    proc = subprocess.run(
        ["systemctl", "restart", "fan-aggressor.service"],
        capture_output=True, text=True
    )
    if proc.returncode != 0:
        sys.stderr.write(proc.stderr or "Failed to restart service")
        sys.exit(1)


def apply_cpu_power():
    for p in [str(Path(__file__).parent), "/usr/local/lib/fan-aggressor"]:
        if p not in sys.path:
            sys.path.insert(0, p)

    from cpu_power import (
        set_governor, set_turbo, set_epp,
        set_rapl_pl1, set_rapl_pl2, set_cpu_max_freq,
    )

    params = json.loads(sys.stdin.read())
    gov = params.get("governor")
    turbo = params.get("turbo", True)
    epp = params.get("epp")
    pp = params.get("platform_profile")

    if gov:
        set_governor(gov)
    set_turbo(turbo)
    if epp:
        set_epp(epp, platform_profile=pp)

    pl1 = params.get("pl1_watts")
    if pl1 is not None:
        set_rapl_pl1(int(pl1))

    pl2 = params.get("pl2_watts")
    if pl2 is not None:
        set_rapl_pl2(int(pl2))

    max_freq = params.get("max_freq_mhz")
    if max_freq is not None:
        set_cpu_max_freq(int(max_freq))


def main():
    if os.getuid() != 0:
        sys.stderr.write("This helper must be run as root (via pkexec)\n")
        sys.exit(1)

    if len(sys.argv) < 2 or sys.argv[1] not in ALLOWED_ACTIONS:
        sys.stderr.write(f"Usage: {sys.argv[0]} <{'|'.join(sorted(ALLOWED_ACTIONS))}>\n")
        sys.exit(1)

    action = sys.argv[1]

    if action == "save-config":
        save_config()
    elif action == "restart-service":
        restart_service()
    elif action == "apply-cpu-power":
        apply_cpu_power()


if __name__ == "__main__":
    main()