cleanup_stale_daemons.py

#!/usr/bin/env python3

import argparse
import os
import plistlib
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path


COLOR = {
    "reset": "\033[0m",
    "bold": "\033[1m",
    "dim": "\033[2m",
    "red": "\033[31m",
    "green": "\033[32m",
    "yellow": "\033[33m",
    "blue": "\033[34m",
    "cyan": "\033[36m",
}

STALE_SERVICES = (
    "com.koekeishiya.yabai",
    "com.koekeishiya.skhd",
)


@dataclass(frozen=True)
class Service:
    label: str
    plist_path: Path
    service_target: str


def supports_color() -> bool:
    return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None


def paint(text: str, *styles: str) -> str:
    if not supports_color():
        return text

    prefix = "".join(COLOR[style] for style in styles)
    return f"{prefix}{text}{COLOR['reset']}"


def status(mark: str, message: str, *styles: str, file=sys.stdout) -> None:
    print(f"{paint(mark, *styles)} {message}", file=file)


def run(argv: list[str], *, apply: bool) -> int:
    command = " ".join(argv)
    if not apply:
        status("DRY", command, "cyan")
        return 0

    status("RUN", command, "blue")
    completed = subprocess.run(argv, check=False)
    return completed.returncode


def read_label(plist_path: Path) -> str | None:
    try:
        with plist_path.open("rb") as handle:
            plist = plistlib.load(handle)
    except (OSError, plistlib.InvalidFileException) as error:
        status(
            "WARN", f"could not read {plist_path}: {error}", "yellow", file=sys.stderr
        )
        return None

    label = plist.get("Label")
    return label if isinstance(label, str) else None


def backup_path(label: str, backup_dir: Path) -> Path:
    candidate = backup_dir / f"{label}.plist.old"
    if not candidate.exists():
        return candidate

    index = 1
    while True:
        candidate = backup_dir / f"{label}.plist.old.{index}"
        if not candidate.exists():
            return candidate
        index += 1


def move_plist(plist_path: Path, destination: Path, *, apply: bool) -> int:
    if not apply:
        status("DRY", f"mv {plist_path} {destination}", "cyan")
        return 0

    status("RUN", f"mv {plist_path} {destination}", "blue")
    try:
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(str(plist_path), str(destination))
    except OSError as error:
        status("ERR", f"could not move {plist_path}: {error}", "red", file=sys.stderr)
        return 1

    return 0


def cleanup_service(service: Service, backup_dir: Path, *, apply: bool) -> int:
    if not service.plist_path.exists():
        status("OK", f"{service.label}: no stale plist found", "green")
        print(f"   {paint(str(service.plist_path), 'dim')}")
        return 0

    plist_label = read_label(service.plist_path)
    if plist_label and plist_label != service.label:
        status(
            "ERR",
            f"{service.plist_path} has label {plist_label!r}, expected {service.label!r}",
            "red",
            file=sys.stderr,
        )
        return 1

    status("CLEAN", f"{service.label}: stale LaunchAgent found", "yellow", "bold")
    result = 0

    domain_target = f"gui/{os.getuid()}"
    bootout = run(
        ["/bin/launchctl", "bootout", domain_target, str(service.plist_path)],
        apply=apply,
    )
    if bootout != 0:
        status(
            "WARN",
            f"bootout failed for {service.label}; continuing with disable",
            "yellow",
            file=sys.stderr,
        )

    disable = run(["/bin/launchctl", "disable", service.service_target], apply=apply)
    if disable != 0:
        status("WARN", f"disable failed for {service.label}", "yellow", file=sys.stderr)
        result = 1

    destination = backup_path(service.label, backup_dir)
    move_result = move_plist(service.plist_path, destination, apply=apply)
    return result or move_result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Clean stale koekeishiya yabai/skhd LaunchAgents after migrating to "
            "asmvik service labels."
        )
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="perform the cleanup; without this flag only print the actions",
    )
    parser.add_argument(
        "--backup-dir",
        type=Path,
        default=Path("/tmp"),
        help="directory where stale plists are moved (default: /tmp)",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    launch_agents_dir = Path.home() / "Library" / "LaunchAgents"
    services = [
        Service(
            label=label,
            plist_path=launch_agents_dir / f"{label}.plist",
            service_target=f"gui/{os.getuid()}/{label}",
        )
        for label in STALE_SERVICES
    ]

    if not args.apply:
        status("INFO", "dry-run mode; pass --apply to change launchd state", "cyan")

    result = 0
    for service in services:
        result = cleanup_service(service, args.backup_dir, apply=args.apply) or result

    return result


if __name__ == "__main__":
    raise SystemExit(main())