kiddo/backend/actions.py

import logging
import shutil
import subprocess
import time
from typing import List, Optional

from backend.settings import get_settings

logger = logging.getLogger(__name__)


class ActionError(Exception):
    pass


def _run(cmd: List[str], env: Optional[dict] = None) -> subprocess.CompletedProcess:
    """Run a command respecting dry-run mode."""
    settings = get_settings()
    if settings.dry_run:
        logger.info("[dry-run] %s", " ".join(cmd))
        return subprocess.CompletedProcess(cmd, 0, stdout="[dry-run]", stderr="")
    return subprocess.run(cmd, check=True, capture_output=True, text=True, env=env)


def _user_exists(user: str) -> bool:
    try:
        return _run(["id", user]).returncode == 0
    except subprocess.CalledProcessError:
        return False


def _get_uid(user: str) -> Optional[str]:
    try:
        result = _run(["id", "-u", user])
        return result.stdout.strip()
    except subprocess.CalledProcessError:
        return None


def _is_logged_in(user: str) -> bool:
    try:
        result = _run(["who"])
        for line in result.stdout.splitlines():
            if line.strip() == "[dry-run]":
                continue
            if line.startswith(f"{user} "):
                return True
    except subprocess.CalledProcessError:
        logger.warning("Could not determine login status for %s", user)
    return False


def _play_sound_if_available() -> Optional[str]:
    settings = get_settings()
    if not settings.sound_player:
        return "sound player not configured"
    if not shutil.which(settings.sound_player):
        return f"{settings.sound_player} not found"
    try:
        _run([settings.sound_player, settings.sound_file])
    except subprocess.CalledProcessError as exc:
        logger.warning("Sound failed: %s", exc)
        return "sound failed"
    return None


def _notify_user(user: str, message: str) -> Optional[str]:
    settings = get_settings()
    notify = shutil.which(settings.notify_send_path)
    if not notify:
        return "notify-send not available"
    uid = _get_uid(user)
    if not uid:
        return "could not resolve uid for notification"

    env = {
        "DISPLAY": ":0",
        "DBUS_SESSION_BUS_ADDRESS": f"unix:path=/run/user/{uid}/bus",
    }
    cmd = [
        "sudo",
        "-u",
        user,
        notify,
        "Safe Kiddo",
        message,
        "-t",
        str(settings.notify_timeout * 1000),
    ]
    try:
        _run(cmd, env=env)
    except subprocess.CalledProcessError as exc:
        logger.warning("Notification failed: %s", exc)
        return "notification failed"
    return None


def disable_user(
    user: str,
    countdown: Optional[int] = None,
    sound: Optional[bool] = None,
    message: Optional[str] = None,
) -> List[str]:
    settings = get_settings()
    if not _user_exists(user):
        raise ActionError(f"user '{user}' not found")

    countdown_seconds = settings.default_countdown if countdown is None else countdown
    play_sound = settings.default_sound if sound is None else sound
    steps: List[str] = []

    _run(["sudo", "usermod", "-L", user])
    steps.append("account locked")

    logged_in = _is_logged_in(user)
    if logged_in:
        steps.append("user is logged in")
        info_message = message or f"Shutdown in {countdown_seconds} seconds. Save your work."

        notify_result = _notify_user(user, info_message)
        if notify_result:
            steps.append(notify_result)
        if play_sound:
            sound_result = _play_sound_if_available()
            if sound_result:
                steps.append(sound_result)

        if countdown_seconds > 0 and not settings.dry_run:
            for remaining in range(countdown_seconds, 0, -1):
                time.sleep(1)
                if remaining % 10 == 0 or remaining <= 5:
                    _notify_user(user, f"Shutdown in {remaining} seconds.")
                    if play_sound:
                        _play_sound_if_available()

    try:
        _run(["sudo", "pkill", "-KILL", "-u", user])
        steps.append("sessions terminated")
    except subprocess.CalledProcessError as exc:
        if exc.returncode == 1:
            steps.append("no sessions to terminate")
        else:
            raise

    if logged_in:
        _run(["sudo", "shutdown", "now"])
        steps.append("system shutdown triggered")
    else:
        steps.append("no active session; shutdown skipped")

    return steps


def enable_user(user: str) -> List[str]:
    if not _user_exists(user):
        raise ActionError(f"user '{user}' not found")
    _run(["sudo", "usermod", "-U", user])
    return ["account unlocked"]


def list_logged_in_users() -> List[str]:
    try:
        result = _run(["who"])
        users: List[str] = []
        for line in result.stdout.splitlines():
            if line.strip() in ("", "[dry-run]"):
                continue
            users.append(line.split()[0])
        return users
    except subprocess.CalledProcessError:
        return []