utils.py

"""Shared utilities for Pi Bootstrap deployment and revert scripts."""

import subprocess
import sys


# ANSI color codes
class Colors:
    CYAN = '\033[96m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    GREEN = '\033[92m'
    BOLD = '\033[1m'
    DIM = '\033[2m'
    RESET = '\033[0m'


def print_colored(message, color=''):
    """Print with ANSI color codes."""
    print(f"{color}{message}{Colors.RESET}")


def print_banner(title, subtitle=''):
    """Print a simple banner."""
    width = 60
    print("\n" + "=" * width)
    print(f"{Colors.BOLD}{title}{Colors.RESET}")
    if subtitle:
        print(f"{Colors.DIM}{subtitle}{Colors.RESET}")
    print("=" * width + "\n")


def run_command(cmd, shell=False, interactive=False):
    """Runs a shell command. If interactive, shows output; otherwise captures it."""
    try:
        if interactive:
            # For interactive commands, we just let them run and see the output
            subprocess.check_call(cmd, shell=shell)
        else:
            # For non-interactive, we capture output to hide it unless there's an error
            result = subprocess.run(cmd, shell=shell, capture_output=True, text=True)
            result.check_returncode()
    except subprocess.CalledProcessError as e:
        print_colored(f"Command failed: {' '.join(cmd) if isinstance(cmd, list) else cmd}", Colors.RED)
        # In non-interactive mode, result is available and has stderr
        if not interactive and 'result' in locals() and result.stderr:
            print_colored(result.stderr.strip(), Colors.RED)
        # For interactive mode, the error is already visible to the user
        sys.exit(e.returncode)


def has_ssh_keys(target, key_path=None):
    """Checks if passwordless SSH is available."""
    try:
        cmd = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=3"]
        if key_path:
            cmd.extend(["-i", key_path])
        cmd.extend([target, "exit"])
        subprocess.check_call(
            cmd,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
        return True
    except subprocess.CalledProcessError:
        return False