mpssh

#!/usr/bin/env -S uv run --script
# vi: ft=python
# /// script
# requires-python = ">=3.7"
# dependencies = [
#     "colorama",
# ]
# ///

import json
import asyncio
from typing import Optional
from datetime import datetime
from colorama import Fore, Style, init
import argparse
import sys
from pathlib import Path
import logging
from dataclasses import dataclass

init()


@dataclass
class SSHResult:
    returncode: int
    stdout: str
    stderr: str
    host: str
    error: Optional[str] = None


def setup_logging(raw: bool = False) -> None:
    class ColoredFormatter(logging.Formatter):
        def format(self, record):
            if hasattr(record, "thread_color"):
                if record.levelno == logging.ERROR:
                    record.msg = f"{record.thread_color}[{now()}] {record.host}:{Fore.RED} {record.msg}{Style.RESET_ALL}"
                else:
                    record.msg = f"{record.thread_color}[{now()}] {record.host}: {record.msg}{Style.RESET_ALL}"
            else:
                if not raw and record.levelno == logging.ERROR:
                    record.msg = f"{Fore.RED}[{now()}] {record.msg}{Style.RESET_ALL}"
            return super().format(record)

    # Create handlers for stdout and stderr
    stdout_handler = logging.StreamHandler(sys.stdout)
    stderr_handler = logging.StreamHandler(sys.stderr)

    if raw:
        formatter = logging.Formatter("%(host)s: %(message)s")
        stdout_handler.setFormatter(formatter)
        stderr_handler.setFormatter(formatter)
    else:
        colored_formatter = ColoredFormatter("%(message)s")
        stdout_handler.setFormatter(colored_formatter)
        stderr_handler.setFormatter(colored_formatter)

    # Set level filters
    stdout_handler.addFilter(lambda record: record.levelno < logging.ERROR)
    stderr_handler.setLevel(logging.ERROR)

    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    # Remove any existing handlers
    for hdlr in logger.handlers[:]:
        logger.removeHandler(hdlr)
    logger.addHandler(stdout_handler)
    logger.addHandler(stderr_handler)


def now() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]


def escape_shell_command(command: str) -> str:
    """Escape special shell characters in the command."""
    return command.replace("'", "'\\''")


async def ssh_connect_and_run_command(
    host: str, command: str, thread_color: str, raw: bool = False, timeout: int = 30, stdin_data: Optional[bytes] = None
) -> SSHResult:
    try:
        escaped_command = escape_shell_command(command)
        ssh_command = f"ssh -o ConnectTimeout={timeout} {host} 'exec \"$SHELL\" -l -c \"{escaped_command}\"'"

        process = await asyncio.create_subprocess_shell(
            ssh_command, 
            stdout=asyncio.subprocess.PIPE, 
            stderr=asyncio.subprocess.PIPE,
            stdin=asyncio.subprocess.PIPE if stdin_data is not None else None
        )

        stdout, stderr = (x.decode("utf-8") for x in await process.communicate(input=stdin_data))

        if stdout:
            for line in stdout.splitlines():
                if raw:
                    logging.info(line.strip(), extra={"host": host})
                else:
                    logging.info(
                        line.strip(), extra={"host": host, "thread_color": thread_color}
                    )

        if stderr:
            if raw:
                logging.error(stderr.strip(), extra={"host": host})
            else:
                logging.error(
                    stderr.strip(), extra={"host": host, "thread_color": thread_color}
                )

        return SSHResult(
            returncode=process.returncode or 0, stdout=stdout, stderr=stderr, host=host
        )

    except Exception as e:
        error_msg = f"Failed to connect to {host}: {str(e)}"
        logging.error(error_msg)
        return SSHResult(
            returncode=-1, stdout="", stderr=str(e), host=host, error=error_msg
        )


def get_hostgroups_help_text(json_file: str) -> str:
    """Get available hostgroups for help text, with warning if config not found."""
    try:
        json_path = Path(json_file)
        if not json_path.exists():
            return f"\nWARNING: Config file not found at {json_file}"
        
        with open(json_path, "r") as file:
            data = json.load(file)
            hostgroups = list(data.keys())
            if hostgroups:
                return f"\nAvailable hostgroups: {', '.join(sorted(hostgroups))}"
            else:
                return f"\nWARNING: No hostgroups found in {json_file}"
    except json.JSONDecodeError:
        return f"\nWARNING: Invalid JSON in config file {json_file}"
    except Exception:
        return f"\nWARNING: Could not read config file {json_file}"


class CustomArgumentParser(argparse.ArgumentParser):
    def error(self, message):
        # Show help with hostgroups when arguments are missing
        self.print_help()
        sys.stderr.write(f"\nerror: {message}\n")
        sys.exit(2)


async def main() -> int:
    # Set default config path
    default_json_file = str(Path.home() / ".config" / "ssh_hosts.json")
    
    parser = CustomArgumentParser(
        description="Run a command on multiple hosts via SSH.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=get_hostgroups_help_text(default_json_file)
    )
    parser.add_argument(
        "--raw", action="store_true", help="Print raw output without any formatting"
    )
    parser.add_argument(
        "--timeout", type=int, default=15, help="SSH connection timeout in seconds"
    )
    parser.add_argument(
        "--json-file",
        type=str,
        help="The JSON file containing the list of hosts",
        default=default_json_file,
    )
    parser.add_argument(
        "--stdin", action="store_true", help="Read stdin and pass it to the remote command"
    )
    parser.add_argument(
        "hostgroup", type=str, help="The hostgroup to run the command on"
    )
    parser.add_argument(
        "command", nargs=argparse.REMAINDER, help="The command to run on each host"
    )
    
    # Handle custom json-file for help text
    if "--json-file" in sys.argv:
        try:
            json_file_idx = sys.argv.index("--json-file") + 1
            if json_file_idx < len(sys.argv) and not sys.argv[json_file_idx].startswith("-"):
                custom_json_file = sys.argv[json_file_idx]
                parser.epilog = get_hostgroups_help_text(custom_json_file)
        except (IndexError, ValueError):
            pass
    
    args = parser.parse_args()

    setup_logging(args.raw)

    try:
        json_path = Path(args.json_file)
        if not json_path.exists():
            raise FileNotFoundError(f"Config file not found at {args.json_file}")

        with open(json_path, "r") as file:
            data = json.load(file)
    except json.JSONDecodeError as e:
        logging.error(f"Invalid JSON in config file {args.json_file}")
        return 1
    except Exception as e:
        logging.error(f"Could not read config file {args.json_file}")
        return 1

    hostgroup = data.get(args.hostgroup)
    if not hostgroup:
        logging.error(f"Hostgroup '{args.hostgroup}' not found in JSON file")
        return 1

    stdin_data = None
    if args.stdin:
        stdin_data = sys.stdin.buffer.read()
    else:
        if not sys.stdin.isatty():
            sys.stdin.read()


    thread_colors = [Fore.GREEN, Fore.YELLOW, Fore.BLUE, Fore.MAGENTA, Fore.CYAN]
    tasks = []

    if not args.raw:
        logging.info(
            f"Executing commands on {len(hostgroup)} host{'' if len(hostgroup) == 1 else 's'}..."
        )

    for i, host in enumerate(hostgroup):
        thread_color = thread_colors[i % len(thread_colors)]
        task = asyncio.create_task(
            ssh_connect_and_run_command(
                host, " ".join(args.command), thread_color, args.raw, args.timeout, stdin_data
            )
        )
        tasks.append(task)

    results = await asyncio.gather(*tasks)

    if not args.raw:
        logging.info("All commands completed.")

    failed_hosts = [r.host for r in results if r.returncode != 0]
    if failed_hosts:
        if args.raw:
            logging.error(
                f"Failed hosts: {', '.join(failed_hosts)}", extra={"host": "all"}
            )
        else:
            logging.error(
                f"Commands failed on {len(failed_hosts)} host{'' if len(failed_hosts) == 1 else 's'}: hosts: {', '.join(failed_hosts)}",
                extra={"host": "all"},
            )
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))