Multi-Device SSH Command Runner

#!/usr/bin/env python3
"""
Multi-Device SSH Command Runner

READ THE DOCUMENTATION BEFORE YOU RUN THIS:
    https://permitany.com/scripts/ssh-command-runner

It covers what the script does, whether it is safe to run, the
inventory file and credentials it needs, how to read the output,
and what to do when something fails.

    pip install "netmiko>=4.3" "PyYAML>=6.0"
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

try:
    import yaml
except ImportError:
    yaml = None

from netmiko import ConnectHandler
from netmiko.exceptions import NetmikoAuthenticationException, NetmikoTimeoutException
import re
from datetime import datetime, timezone

# ---------------------------------------------------------------------------
# Inventory loading and connection
# ---------------------------------------------------------------------------

NETMIKO_DRIVER = {
    "ios": "cisco_ios",
    "iosxe": "cisco_ios",
    "iosxr": "cisco_xr",
    "nxos": "cisco_nxos",
    "junos": "juniper_junos",
    "eos": "arista_eos",
    "panos": "paloalto_panos",
    "fortios": "fortinet",
}


def load_inventory(path: str) -> list[dict]:
    """
    Read devices.yaml (or .json) and return a list of device dicts.

    Applies the `defaults:` block to every device, validates the `os` field,
    and fails with a readable message rather than a stack trace.
    """
    p = Path(path)
    if not p.exists():
        raise SystemExit(
            f"Inventory file not found: {path}\n"
            f"Create it with:\n\n"
            f"  devices:\n"
            f"    - host: 192.168.1.1\n"
            f"      os: ios\n"
        )

    raw = p.read_text(encoding="utf-8")
    try:
        if p.suffix.lower() in (".yaml", ".yml"):
            if yaml is None:
                raise SystemExit("PyYAML is not installed. Run:  pip install PyYAML")
            data = yaml.safe_load(raw)
        else:
            data = json.loads(raw)
    except SystemExit:
        raise
    except Exception as exc:
        raise SystemExit(f"Could not parse {path}: {exc}\n"
                         f"Check for tabs (use spaces) and missing spaces after colons.")

    if isinstance(data, dict):
        defaults = data.get("defaults", {}) or {}
        devices = data.get("devices", []) or []
    else:
        defaults, devices = {}, data

    if not devices:
        raise SystemExit(f"No devices found in {path}. The file needs a `devices:` list.")

    out = []
    for i, d in enumerate(devices, 1):
        if not isinstance(d, dict):
            raise SystemExit(f"Device #{i} in {path} is not a mapping. "
                             f"Each entry needs at least `host:` and `os:`.")
        merged = {**defaults, **d}
        if "host" not in merged:
            raise SystemExit(f"Device #{i} in {path} has no `host:` field.")
        os_key = merged.get("os")
        if not os_key:
            raise SystemExit(f"Device '{merged['host']}' has no `os:` field. "
                             f"Valid values: {', '.join(sorted(NETMIKO_DRIVER))}")
        if os_key not in NETMIKO_DRIVER:
            raise SystemExit(
                f"Device '{merged['host']}' has os '{os_key}', which is not supported.\n"
                f"Valid values: {', '.join(sorted(NETMIKO_DRIVER))}\n"
                f"Note: 'cisco' is a vendor, not an OS - use ios, iosxe, iosxr, or nxos."
            )
        merged.setdefault("vrfs", [])
        out.append(merged)
    return out


def credentials() -> tuple[str, str, str]:
    """Read credentials from the environment and fail clearly if missing."""
    user = os.environ.get("NET_USER", "")
    pwd = os.environ.get("NET_PASS", "")
    enable = os.environ.get("NET_ENABLE", "")
    if not user or not pwd:
        raise SystemExit(
            "Set NET_USER and NET_PASS in your environment before running.\n\n"
            "  Linux/macOS:   export NET_USER='admin'\n"
            "                 export NET_PASS='your-password'\n\n"
            "  PowerShell:    $env:NET_USER = 'admin'\n"
            "                 $env:NET_PASS = 'your-password'\n"
        )
    return user, pwd, enable


def connect(dev: dict):
    """Open an SSH session to one device using netmiko."""
    user, pwd, enable = credentials()
    params = {
        "device_type": NETMIKO_DRIVER[dev["os"]],
        "host": dev["host"],
        "username": dev.get("username") or user,
        "password": dev.get("password") or pwd,
        "port": dev.get("port", 22),
        "conn_timeout": 20,
        "fast_cli": False,
    }
    secret = dev.get("secret") or enable
    if secret:
        params["secret"] = secret

    conn = ConnectHandler(**params)
    if secret and dev["os"] in ("ios", "iosxe", "nxos"):
        conn.enable()
    return conn


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------

# Default command per OS when you do not pass --command.
DEFAULT_COMMANDS = {
    "ios": "show version",
    "iosxe": "show version",
    "iosxr": "show version",
    "nxos": "show version",
    "junos": "show version",
    "eos": "show version",
    "panos": "show system info",
    "fortios": "get system status",
}

# This script is READ-ONLY by design. Anything that could enter config mode or
# change device state is refused before a single session is opened. Use
# config-rollback or the bulk-* scripts for changes - they have snapshots,
# verification, and undo. This one has none of that.
WRITE_COMMANDS = re.compile(
    r"^\s*(conf(ig(ure)?)?\b|wr(ite)?\b|copy\b|reload\b|reboot\b|erase\b|delete\b|"
    r"clear\b|no\s|set\b|request\b|execute\b|install\b|rollback\b|commit\b|"
    r"archive\b|format\b|shutdown\b)",
    re.IGNORECASE,
)

# How each OS scopes a show command to a VRF / routing-instance / virtual-router.
VRF_SUFFIX = {
    "ios": "vrf {vrf}",
    "iosxe": "vrf {vrf}",
    "iosxr": "vrf {vrf}",
    "nxos": "vrf {vrf}",
    "eos": "vrf {vrf}",
    "junos": "instance {vrf}",
    "fortios": "vrf {vrf}",
    "panos": "virtual-router {vrf}",
}


def build_command(base: str, os_key: str, vrf: str | None) -> str:
    if not vrf:
        return base
    suffix = VRF_SUFFIX.get(os_key, "")
    return f"{base} {suffix.format(vrf=vrf)}".strip() if suffix else base


def load_command_map(path: str | None, override: str | None) -> dict[str, str]:
    """
    Resolve which command runs on which OS.

    --command wins for every OS. --commands-file lets you run a different
    command per OS in one job, which matters because the same question has
    different syntax on each platform.
    """
    if override:
        return {k: override for k in DEFAULT_COMMANDS}
    if not path:
        return dict(DEFAULT_COMMANDS)

    p = Path(path)
    if not p.exists():
        raise SystemExit(f"Commands file not found: {path}")
    raw = p.read_text(encoding="utf-8")
    try:
        data = (yaml.safe_load(raw) if p.suffix.lower() in (".yaml", ".yml") and yaml
                else json.loads(raw))
    except Exception as exc:
        raise SystemExit(f"Could not parse {path}: {exc}")

    cmds = data.get("commands", data) if isinstance(data, dict) else {}
    if not isinstance(cmds, dict) or not cmds:
        raise SystemExit(
            f"{path} must map an os key to a command, for example:\n\n"
            f"  commands:\n"
            f"    ios: show ip interface brief\n"
            f"    junos: show interfaces terse\n"
            f"    fortios: get system interface physical\n"
        )
    unknown = set(cmds) - set(NETMIKO_DRIVER)
    if unknown:
        raise SystemExit(f"{path} has unknown os key(s): {', '.join(sorted(unknown))}\n"
                         f"Valid: {', '.join(sorted(NETMIKO_DRIVER))}")
    return cmds


# ---------------------------------------------------------------------------
# Execution
# ---------------------------------------------------------------------------

def run_on_device(dev: dict, commands: dict[str, str], args) -> dict:
    host = dev["host"]
    os_key = dev["os"]
    base = commands.get(os_key)
    conn = None

    if not base:
        return {"host": host, "os": os_key, "status": "skip",
                "note": f"no command defined for os '{os_key}'"}

    outputs: dict[str, str] = {}
    try:
        conn = connect(dev)

        outputs["global"] = conn.send_command(base, read_timeout=args.read_timeout)
        for vrf in dev.get("vrfs", []):
            cmd = build_command(base, os_key, vrf)
            if cmd == base:
                continue  # this OS has no VRF form for show commands
            outputs[vrf] = conn.send_command(cmd, read_timeout=args.read_timeout)

        if args.output_dir:
            d = Path(args.output_dir)
            d.mkdir(parents=True, exist_ok=True)
            ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
            body = "\n".join(
                f"! === {host} ({os_key}) [{ctx}] {ts} UTC\n! {build_command(base, os_key, None if ctx == 'global' else ctx)}\n{out}"
                for ctx, out in outputs.items()
            )
            (d / f"{host}_{ts}.txt").write_text(body, encoding="utf-8")

        return {"host": host, "os": os_key, "status": "ok",
                "command": base, "outputs": outputs}

    except NetmikoAuthenticationException as exc:
        return {"host": host, "os": os_key, "status": "error",
                "error": f"authentication failed - check NET_USER/NET_PASS ({exc})"}
    except NetmikoTimeoutException as exc:
        return {"host": host, "os": os_key, "status": "error",
                "error": f"unreachable or SSH timed out ({exc})"}
    except Exception as exc:  # noqa: BLE001 - one device must not abort the run
        return {"host": host, "os": os_key, "status": "error",
                "error": f"{type(exc).__name__}: {exc}"}
    finally:
        if conn:
            try:
                conn.disconnect()
            except Exception:  # noqa: BLE001
                pass


def main() -> int:
    ap = argparse.ArgumentParser(
        description="Run a read-only show command across many devices at once.",
        epilog="Credentials come from NET_USER and NET_PASS. See the docstring at "
               "the top of this file for the full inventory format.",
    )
    ap.add_argument("--inventory", required=True,
                    help="Path to devices.yaml listing which devices to query")
    ap.add_argument("--command", "-c",
                    help="Run this exact command on every device, regardless of OS")
    ap.add_argument("--commands-file",
                    help="YAML/JSON mapping each os key to its own command, for when "
                         "the same question has different syntax per platform")
    ap.add_argument("--limit", help="Only query hosts whose name contains this text")
    ap.add_argument("--os", dest="os_filter",
                    help="Only query devices running this OS, e.g. --os nxos")
    ap.add_argument("--output-dir",
                    help="Also save each device's output to a file in this directory")
    ap.add_argument("--grep",
                    help="Only print output lines matching this regex. The full "
                         "output is still saved when --output-dir is used")
    ap.add_argument("--quiet", action="store_true",
                    help="Print only the status line per device, not the output")
    ap.add_argument("--read-timeout", type=int, default=90,
                    help="Seconds to wait for each command. Default: 90")
    ap.add_argument("--workers", type=int, default=10,
                    help="How many devices to contact at once. Default: 10")
    ap.add_argument("--json-out", help="Write machine-readable results here")
    args = ap.parse_args()

    if args.command and args.commands_file:
        raise SystemExit("Use --command or --commands-file, not both.")

    commands = load_command_map(args.commands_file, args.command)

    # Read-only guard, checked before any connection is opened.
    for os_key, cmd in commands.items():
        if WRITE_COMMANDS.match(cmd):
            raise SystemExit(
                f"Refusing to run '{cmd}' - this script is read-only.\n\n"
                f"It has no backup, no verification, and no undo. For configuration\n"
                f"changes use config-rollback (snapshot + automatic rollback) or one\n"
                f"of the bulk-* scripts, which are built for writes."
            )

    devices = load_inventory(args.inventory)
    if args.limit:
        devices = [d for d in devices if args.limit in d["host"]]
    if args.os_filter:
        devices = [d for d in devices if d["os"] == args.os_filter]
    if not devices:
        raise SystemExit("No device matched the --limit / --os filters.")

    credentials()  # fail fast before opening any connection

    label = args.command or ("per-OS commands from " + args.commands_file
                             if args.commands_file else "OS-specific defaults")
    print(f"Running: {label}")
    print(f"Devices: {len(devices)}")
    print("=" * 70)

    pattern = re.compile(args.grep) if args.grep else None
    results = []

    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(run_on_device, d, commands, args): d["host"]
                   for d in devices}
        for fut in as_completed(futures):
            host = futures[fut]
            try:
                r = fut.result()
            except Exception as exc:  # noqa: BLE001
                r = {"host": host, "os": "?", "status": "error",
                     "error": f"worker crashed: {exc}"}
            results.append(r)

            tag = {"ok": "[OK]", "error": "[FAIL]", "skip": "[SKIP]"}[r["status"]]
            head = f"{tag}  {r['host']} ({r['os']})"

            if r["status"] == "error":
                print(f"\n{head}\n       {r['error']}")
                continue
            if r["status"] == "skip":
                print(f"\n{head}\n       {r['note']}")
                continue

            print(f"\n{head}  $ {r['command']}")
            if args.quiet:
                continue
            for ctx, out in r["outputs"].items():
                if ctx != "global":
                    print(f"  --- VRF: {ctx} ---")
                lines = out.splitlines()
                if pattern:
                    lines = [ln for ln in lines if pattern.search(ln)]
                    if not lines:
                        print("       (no line matched --grep)")
                        continue
                for ln in lines:
                    print(f"       {ln}")

    if args.json_out:
        Path(args.json_out).write_text(json.dumps(results, indent=2), encoding="utf-8")

    ok = sum(1 for r in results if r["status"] == "ok")
    failed = sum(1 for r in results if r["status"] == "error")
    skipped = sum(1 for r in results if r["status"] == "skip")
    print("\n" + "=" * 70)
    print(f"Done. {ok} ok, {skipped} skipped, {failed} failed (of {len(devices)}).")
    return 1 if failed else 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\nInterrupted.", file=sys.stderr)
        sys.exit(130)