pyMC_Repeater/repeater/service_utils.py

"""
Service management utilities for pyMC Repeater.
Provides functions for service control operations like restart.
"""

import logging
import os
import subprocess
import threading
import time
from typing import Dict, Optional, Tuple

logger = logging.getLogger("ServiceUtils")
INIT_SCRIPT = "/etc/init.d/S80pymc-repeater"
BUILDROOT_METADATA_PATH = "/etc/pymc-image-build-id"
_CONTAINER_RESTART_DELAY_SECONDS = 1.0


def is_buildroot() -> bool:
    if os.path.exists(BUILDROOT_METADATA_PATH):
        return True
    if os.path.exists("/etc/os-release"):
        try:
            with open("/etc/os-release", "r", encoding="utf-8") as handle:
                return any(line.strip() == "ID=buildroot" for line in handle)
        except OSError:
            return False
    return False


def get_buildroot_image_info() -> Dict[str, str]:
    info: Dict[str, str] = {}

    try:
        with open(BUILDROOT_METADATA_PATH, "r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line or "=" not in line:
                    continue
                key, value = line.split("=", 1)
                info[key.strip()] = value.strip()
    except OSError:
        return {}

    return info


def get_buildroot_image_version() -> Optional[str]:
    return get_buildroot_image_info().get("image_version")


def is_container() -> bool:
    """Detect common Docker/LXC/containerized environments."""
    if os.path.exists("/.dockerenv") or os.environ.get("container"):
        return True

    try:
        with open("/proc/1/environ", "rb") as handle:
            if b"container=" in handle.read():
                return True
    except (OSError, PermissionError):
        pass

    try:
        with open("/proc/1/cgroup", "r", encoding="utf-8") as handle:
            cgroup_data = handle.read()
            if any(token in cgroup_data for token in ("docker", "containerd", "kubepods", "lxc")):
                return True
    except OSError:
        pass

    return os.path.exists("/run/host/container-manager")


def _schedule_container_exit(delay_seconds: float = _CONTAINER_RESTART_DELAY_SECONDS) -> None:
    """Exit the current process shortly after returning success to the caller."""

    def _exit_process() -> None:
        time.sleep(delay_seconds)
        logger.warning("Exiting repeater process to trigger container restart")
        os._exit(0)

    threading.Thread(target=_exit_process, name="container-restart-exit", daemon=True).start()


def get_container_restart_message() -> str:
    """Return the user-facing restart message for containerized installs."""
    return (
        "Container restart initiated. "
        "If you are running pyMC Repeater via Docker or Home Assistant, pull or rebuild "
        "a newer image for packaged image updates to take effect."
    )


def restart_service() -> Tuple[bool, str]:
    """
    Restart the pymc-repeater service.
    
    On Buildroot/Luckfox, use the shipped init script directly.
    On systemd hosts, try polkit-based restart first (plain systemctl), then
    fall back to sudo-based restart (requires sudoers.d rule installed by
    manage.sh).
    
    Returns:
        Tuple[bool, str]: (success, message)
    """
    if is_container():
        _schedule_container_exit()
        logger.info("Container environment detected; scheduled process exit for container restart")
        return True, get_container_restart_message()

    if is_buildroot():
        if not os.path.exists(INIT_SCRIPT):
            logger.error("Buildroot init script not found: %s", INIT_SCRIPT)
            return False, f"init script not found: {INIT_SCRIPT}"

        try:
            subprocess.Popen(
                ["/bin/sh", "-c", f"sleep 1; exec {INIT_SCRIPT} restart >/dev/null 2>&1"],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                stdin=subprocess.DEVNULL,
                start_new_session=True,
            )
            logger.info("Service restart scheduled via Buildroot init script")
            return True, "Service restart initiated"
        except Exception as exc:
            logger.error(f"Buildroot restart failed: {exc}")
            return False, f"Restart failed: {exc}"

    # Try polkit-based restart first (works on bare metal / VMs with polkit running)
    try:
        result = subprocess.run(
            ["systemctl", "restart", "pymc-repeater"], capture_output=True, text=True, timeout=5
        )

        if result.returncode == 0:
            logger.info("Service restart via polkit succeeded")
            return True, "Service restart initiated"
        
        stderr = result.stderr or ""
        if "Access denied" in stderr or "authorization" in stderr.lower():
            logger.info("Polkit denied restart, trying sudo fallback...")
        else:
            # Some other error, still try sudo
            logger.warning(f"systemctl restart failed ({result.returncode}): {stderr.strip()}")
            
    except subprocess.TimeoutExpired:
        # Timeout likely means it's restarting - that's success
        logger.warning("Service restart command timed out (service may be restarting)")
        return True, "Service restart initiated (timeout - likely restarting)"
    except FileNotFoundError:
        logger.error("systemctl not found")
        return False, "systemctl not available"
    except Exception as e:
        logger.warning(f"Polkit restart attempt failed: {e}")
    
    # Fallback: use sudo (requires /etc/sudoers.d/pymc-repeater rule)
    try:
        result = subprocess.run(
            ['sudo', '--non-interactive', 'systemctl', 'restart', 'pymc-repeater'],
            capture_output=True,
            text=True,
            timeout=5
        )
        
        if result.returncode == 0:
            logger.info("Service restart via sudo succeeded")
            return True, "Service restart initiated"
        else:
            error_msg = result.stderr or "Unknown error"
            logger.error(f"Service restart via sudo failed: {error_msg}")
            return False, f"Restart failed: {error_msg}"
            
    except subprocess.TimeoutExpired:
        logger.warning("Sudo restart timed out (service likely restarting)")
        return True, "Service restart initiated (timeout - likely restarting)"
    except FileNotFoundError:
        logger.error("sudo not found - cannot restart service")
        return False, "Neither polkit nor sudo available for service restart"
    except Exception as e:
        logger.error(f"Error executing sudo restart: {e}")
        return False, f"Restart command failed: {str(e)}"