wrolpi/wrolpi/cmd.py

"""Paths for third-party commands that WROLPi may require."""
import asyncio
import dataclasses
import os
import pathlib
import shlex
import shutil
import signal
import tempfile
from asyncio import CancelledError
from pathlib import Path
from time import time
from typing import Optional

from wrolpi.common import logger, TRACE_LEVEL
from wrolpi.vars import PYTEST, DOCKERIZED

logger = logger.getChild(__name__)


def is_executable(path: Path) -> bool:
    """Returns True if the file is executable."""
    return path.is_file() and os.access(path, os.X_OK)


def which(*possible_paths: str, warn: bool = False) -> Optional[Path]:
    """
    Find an executable in the system $PATH.  If the executable cannot be found in $PATH, then return the first
    executable found in possible_paths.

    Returns None if no executable can be found.
    """
    found = None
    possible_paths = list(map(Path, possible_paths))
    for path in possible_paths:
        if is_executable(path):
            # `path` was an absolute file which is executable.
            found = path
            break
        path = shutil.which(path)
        if path and (path := Path(path).absolute()) and is_executable(path):
            # `path` was a name or relative path.
            found = path
            break

    if not found and warn and not PYTEST and not DOCKERIZED:
        logger.warning(f'Cannot find executable {possible_paths[0]}')
        return
    elif found:

        return found.absolute()


def pid_is_running(pid: int) -> bool:
    """Return True if a process with PID exists."""
    try:
        os.kill(pid, 0)
        return True
    except OSError:
        return False


# Admin
SUDO_BIN = which('sudo', '/usr/bin/sudo')
NMCLI_BIN = which('nmcli', '/usr/bin/nmcli')
CPUFREQ_INFO_BIN = which('cpufreq-info', '/usr/bin/cpufreq-info')
CPUFREQ_SET_BIN = which('cpufreq-set', '/usr/bin/cpufreq-set')

# Files
WGET_BIN = which('wget', '/usr/bin/wget')
GIT_BIN = which('git', '/usr/bin/git')

# Map
BASH_BIN = which('bash', '/bin/bash')

# Archive
SINGLE_FILE_BIN = which('single-file',
                        '/usr/bin/single-file',  # rpi os
                        '/usr/local/bin/single-file',  # debian
                        )
DENO_BIN = which('deno', '/usr/local/bin/deno')


def _find_single_file_deno_script() -> Optional[Path]:
    """Find the Deno entry script of the single-file-cli package.

    npm's `single-file` bin is a symlink to the package's single-file-node.js; the Deno script
    (`single-file`, shebang `deno run`) lives next to it.  single-file-cli is designed for Deno —
    under node 18 it cannot reach the browser CDP on IPv6-enabled hosts and never exits."""
    if not SINGLE_FILE_BIN:
        return None
    script = SINGLE_FILE_BIN.resolve().parent / 'single-file'
    try:
        with script.open('rb') as fh:
            if b'deno run' in fh.readline():
                return script
    except OSError:
        pass
    return None


SINGLE_FILE_DENO_SCRIPT = _find_single_file_deno_script()
CHROMIUM = which('chromium-browser', 'chromium',
                 '/usr/bin/chromium-browser',  # rpi os
                 '/usr/bin/chromium',  # debian
                 )
CHROMEDRIVER_BIN = which('chromedriver',
                         '/usr/bin/chromedriver',  # rpi os / debian
                         )
FIREFOX = which('firefox', 'firefox-esr',
                '/usr/bin/firefox',
                '/usr/bin/firefox-esr',
                )
READABILITY_BIN = which('readability-extractor',
                        '/usr/bin/readability-extractor',  # rpi os
                        '/usr/local/bin/readability-extractor',  # debian
                        )

# Known browser definitions for SingleFile
KNOWN_BROWSERS = {
    'chromium': {
        'paths': ['chromium-browser', 'chromium', '/usr/bin/chromium-browser', '/usr/bin/chromium'],
        'name': 'Chromium',
    },
    'firefox': {
        'paths': ['firefox', 'firefox-esr', '/usr/bin/firefox', '/usr/bin/firefox-esr'],
        'name': 'Firefox',
    },
    'brave': {
        'paths': ['brave-browser', '/usr/bin/brave-browser'],
        'name': 'Brave',
    }
}


def get_installed_browsers() -> list:
    """
    Detect browsers installed on the system that can be used with SingleFile.

    Returns a list of dicts with:
        - key: The browser identifier (e.g., 'chromium', 'firefox')
        - name: Human-readable name (e.g., 'Chromium', 'Firefox')
        - path: Absolute path to the executable
    """
    browsers = []

    for key, info in KNOWN_BROWSERS.items():
        path = which(*info['paths'])
        if path:
            browsers.append({
                'key': key,
                'name': info['name'],
                'path': str(path),
            })

    return browsers


# Videos
YT_DLP_BIN = which(
    '/opt/wrolpi/venv/bin/yt-dlp',  # Use virtual environment location
    'yt-dlp',
    '/usr/local/bin/yt-dlp',  # Location in docker container
)
FFPROBE_BIN = which('ffprobe', '/usr/bin/ffprobe')
FFMPEG_BIN = which('ffmpeg', '/usr/bin/ffmpeg')

# Documents
CATDOC_PATH = which('catdoc')
TEXTUTIL_PATH = which('textutil')


@dataclasses.dataclass
class CommandResult:
    return_code: int
    cancelled: bool
    stdout: bytes
    stderr: bytes
    elapsed: int
    pid: int = None


TESTING_RUN_COMMAND_RESULT = None


async def run_command(cmd: tuple[str | pathlib.Path, ...], cwd: pathlib.Path | str = None,
                      timeout: int = 600, log_command: bool = True,
                      stdout_callback: callable = None, env: dict = None,
                      start_new_session: bool = False, merge_stderr: bool = False) -> CommandResult:
    """Run a shell command, return the results (stdout/stderr/return code).

    When stdout_callback is provided, stdout is piped and each line is passed to the callback
    in real-time.  Otherwise, stdout is written to a temporary file.

    :param log_command: Enable debug logging.
    :param stdout_callback: Called with each decoded stdout line (str). May be None.
    :param env: Environment variables for the subprocess. If None, inherits the current environment.
    :param start_new_session: Run the command in its own session/process group; when it finishes
        (or is killed by timeout/cancel), the whole group is killed.  Use for commands that may
        leave children behind (e.g. single-file's browser).
    :param merge_stderr: Send stderr to stdout (and so to `stdout_callback`); `stderr` will be empty.  Use for
        commands that write progress to stderr (e.g. git).
    """
    if not isinstance(cmd, (list, tuple)):
        raise RuntimeError('Command must be a list or tuple')

    if PYTEST and TESTING_RUN_COMMAND_RESULT:
        logger.debug('run_command: returning mock result')
        return TESTING_RUN_COMMAND_RESULT()  # call the mock

    streaming = stdout_callback is not None

    with tempfile.NamedTemporaryFile() as stderr_fh:
        stderr_file = pathlib.Path(stderr_fh.name)
        cmd = tuple(str(i) for i in cmd)
        cmd_str = ' '.join(shlex.quote(i) for i in cmd)
        if log_command:
            logger.info(f'run_command: running ({timeout=}): {cmd_str}')
        elif __debug__ and logger.isEnabledFor(TRACE_LEVEL):
            logger.trace(f'run_command: running ({timeout=}): {cmd_str}')

        stdout_fh = None
        stdout_file = None
        reader_task = None
        stdout_buf = None
        pid = None

        if streaming:
            stdout_arg = asyncio.subprocess.PIPE
            stdout_buf = bytearray()
        else:
            stdout_fh = tempfile.NamedTemporaryFile()
            stdout_file = pathlib.Path(stdout_fh.name)
            stdout_arg = stdout_fh

        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=stdout_arg,
                stderr=asyncio.subprocess.STDOUT if merge_stderr else stderr_fh,
                cwd=cwd,
                env=env,
                start_new_session=start_new_session,
            )
            pid = proc.pid

            if streaming:
                async def _read_stdout():
                    buf = ''
                    while True:
                        chunk = await proc.stdout.read(4096)
                        if not chunk:
                            # Flush remaining buffer.
                            if buf.strip():
                                try:
                                    stdout_callback(buf.strip())
                                except Exception as e:
                                    logger.warning(f'run_command: stdout_callback error: {e}')
                            break
                        stdout_buf.extend(chunk)
                        buf += chunk.decode(errors='replace')
                        # Split on both \n and \r to handle progress bars.
                        while '\n' in buf or '\r' in buf:
                            # Find the earliest delimiter.
                            ni = buf.find('\n')
                            ri = buf.find('\r')
                            if ni == -1:
                                idx, skip = ri, 1
                            elif ri == -1:
                                idx, skip = ni, 1
                            else:
                                idx, skip = min(ni, ri), 1
                            line = buf[:idx]
                            buf = buf[idx + skip:]
                            if line.strip():
                                try:
                                    stdout_callback(line.strip())
                                except Exception as e:
                                    logger.warning(f'run_command: stdout_callback error: {e}')

                reader_task = asyncio.create_task(_read_stdout())

            cancelled = False
            start = time()
            try:
                while True:
                    elapsed = int(time() - start)
                    if timeout and elapsed >= timeout:
                        logger.warning(f'run_command: timeout exceeded, killing... {cmd=}')
                        proc.kill()
                        await proc.wait()
                        break

                    try:
                        await asyncio.wait_for(proc.wait(), timeout=1)
                        break
                    except asyncio.TimeoutError:
                        continue
            except CancelledError as e:
                logger.warning(f'run_command: cancelled, killing... {cmd=}', exc_info=e)
                cancelled = True
                proc.kill()
                await proc.wait()

            if reader_task:
                # Drain remaining stdout.  Cancelling here races the last lines on a fast process
                # (CI often loses "line three" because wait() returns before the reader finishes).
                try:
                    await reader_task
                except (CancelledError, asyncio.CancelledError):
                    pass

            elapsed = int(time() - start)
            if streaming:
                stdout = bytes(stdout_buf)
            else:
                stdout = stdout_file.read_bytes() or b''
            stderr = stderr_file.read_bytes() or b''
            if log_command:
                logger.debug(
                    f'run_command: finished ({elapsed=}s) with stdout={len(stdout)} stderr={len(stderr)}: {cmd_str=}')
            elif __debug__ and logger.isEnabledFor(TRACE_LEVEL):
                logger.trace(
                    f'run_command: finished ({elapsed=}s) with stdout={len(stdout)} stderr={len(stderr)}: {cmd_str=}')
            return CommandResult(
                return_code=proc.returncode,
                cancelled=cancelled,
                stdout=stdout,
                stderr=stderr,
                elapsed=elapsed,
                pid=pid,
            )
        finally:
            if start_new_session and pid:
                # The command was its own session/group leader; kill anything it left behind
                # (e.g. the browser single-file spawns but fails to close).
                try:
                    os.killpg(pid, signal.SIGKILL)
                    logger.info(f'run_command: killed lingering process group {pid}')
                except ProcessLookupError:
                    pass  # Nothing left behind; the common case.
                except PermissionError as e:
                    logger.error(f'run_command: failed to kill process group {pid}', exc_info=e)
            if stdout_fh:
                stdout_fh.close()