Safe Config Push with Automatic Rollback

#!/usr/bin/env python3
"""
Safe Config Push with Automatic Rollback

READ THE DOCUMENTATION BEFORE YOU RUN THIS:
    https://permitany.com/scripts/config-rollback

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 difflib
import json
import logging
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

try:
    import yaml
except ImportError:
    yaml = None

from netmiko import ConnectHandler
from netmiko.exceptions import (
    NetmikoAuthenticationException,
    NetmikoTimeoutException,
    ReadTimeout,
)

LOG = logging.getLogger("rollback")

# --------------------------------------------------------------------------
# Platform matrix
# --------------------------------------------------------------------------

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",
}

SHOW_RUN = {
    "ios": "show running-config",
    "iosxe": "show running-config",
    "iosxr": "show running-config",
    "nxos": "show running-config",
    "junos": "show configuration | display set",
    "eos": "show running-config",
    "panos": "show config running",
    "fortios": "show full-configuration",
}

SAVE_CMD = {
    "ios": "write memory",
    "iosxe": "write memory",
    "iosxr": None,          # commit writes
    "nxos": "copy running-config startup-config",
    "junos": None,          # commit writes
    "eos": "write memory",
    "panos": None,          # commit writes
    "fortios": None,        # auto-saves
}

# Platforms whose device reverts on its own if we never confirm.
NATIVE_CONFIRM = {"iosxr", "junos", "eos"}


# --------------------------------------------------------------------------
# Data model
# --------------------------------------------------------------------------

@dataclass
class Device:
    host: str
    os: str
    username: str = ""
    password: str = ""
    secret: str = ""
    port: int = 22
    vrfs: list[str] = field(default_factory=list)
    # Assertions evaluated before and after the change. Each is a command whose
    # output must stay equal (or satisfy `expect`) for the change to be kept.
    checks: list[dict[str, str]] = field(default_factory=list)

    @property
    def driver(self) -> str:
        if self.os not in NETMIKO_DRIVER:
            raise ValueError(f"{self.host}: unsupported os '{self.os}'")
        return NETMIKO_DRIVER[self.os]


@dataclass
class Result:
    host: str
    os: str
    status: str = "pending"      # ok | rolled_back | failed | skipped | dry_run
    diff_lines: int = 0
    snapshot: str = ""
    message: str = ""
    checks_before: dict[str, str] = field(default_factory=dict)
    checks_after: dict[str, str] = field(default_factory=dict)
    elapsed_s: float = 0.0


# --------------------------------------------------------------------------
# Connection
# --------------------------------------------------------------------------

def connect(dev: Device, read_timeout: int = 90):
    params = {
        "device_type": dev.driver,
        "host": dev.host,
        "username": dev.username or os.environ.get("NET_USER", ""),
        "password": dev.password or os.environ.get("NET_PASS", ""),
        "port": dev.port,
        "conn_timeout": 20,
        "fast_cli": False,
    }
    secret = dev.secret or os.environ.get("NET_ENABLE", "")
    if secret:
        params["secret"] = secret

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


def verify_privilege(conn, dev: Device) -> None:
    """Fail fast if we cannot configure, rather than halfway through a push."""
    if dev.os in ("ios", "iosxe"):
        out = conn.send_command("show privilege")
        if "15" not in out:
            raise PermissionError(f"{dev.host}: privilege level 15 required, got: {out.strip()}")
    elif dev.os == "fortios":
        out = conn.send_command("get system status")
        if "Current virtual domain" not in out and "Version" not in out:
            raise PermissionError(f"{dev.host}: unexpected response, check account scope")


# --------------------------------------------------------------------------
# Snapshot / restore
# --------------------------------------------------------------------------

def fetch_config(conn, dev: Device) -> str:
    cmd = SHOW_RUN[dev.os]
    if dev.os == "fortios":
        conn.send_command("config global", expect_string=r"#")
    raw = conn.send_command(cmd, read_timeout=180)
    return normalise(raw, dev.os)


def normalise(text: str, os_key: str) -> str:
    """Strip volatile lines so diffs reflect real change, not timestamps."""
    noise = (
        "Building configuration",
        "Current configuration",
        "! Last configuration change",
        "! NVRAM config last updated",
        "Time: ",
        "## Last commit:",
        "!Time:",
    )
    out = []
    for line in text.splitlines():
        s = line.strip()
        if not s:
            continue
        if any(s.startswith(n) for n in noise):
            continue
        out.append(line.rstrip())
    return "\n".join(out)


def save_snapshot(cfg: str, dev: Device, stamp: str, root: Path) -> Path:
    d = root / dev.host
    d.mkdir(parents=True, exist_ok=True)
    p = d / f"{stamp}.cfg"
    p.write_text(cfg, encoding="utf-8")
    return p


def onbox_snapshot(conn, dev: Device, tag: str) -> str | None:
    """Create a device-side restore point. Returns its identifier, or None."""
    try:
        if dev.os == "iosxe":
            conn.send_command(f"archive config", read_timeout=60)
            out = conn.send_command_timing(f"copy running-config flash:{tag}.cfg")
            if "?" in out or "filename" in out.lower():
                out += conn.send_command_timing("\n")
            return f"flash:{tag}.cfg"
        if dev.os == "nxos":
            conn.send_command(f"checkpoint {tag}", read_timeout=120)
            return tag
        if dev.os == "fortios":
            # FortiOS keeps an automatic revision table; record the current id
            out = conn.send_command("execute revision list config")
            return out.strip().splitlines()[-1].split()[0] if out.strip() else None
    except Exception as exc:  # noqa: BLE001 - snapshot is best-effort
        LOG.warning("%s: on-box snapshot failed (%s); disk snapshot still taken", dev.host, exc)
    return None


# --------------------------------------------------------------------------
# Config push, per platform
# --------------------------------------------------------------------------

def push_with_confirm(conn, dev: Device, lines: list[str], confirm_min: int) -> str:
    """
    Apply config inside a window the device will undo by itself.
    Returns a token the caller passes to `confirm()` on success.
    """
    os_key = dev.os

    if os_key == "iosxr":
        conn.send_config_set(lines, exit_config_mode=False)
        out = conn.send_command_timing(f"commit confirmed {confirm_min} minutes")
        if "Failed" in out or "error" in out.lower():
            conn.send_command_timing("abort")
            raise RuntimeError(f"commit confirmed rejected: {out.strip()[:300]}")
        conn.exit_config_mode()
        return "iosxr-confirm"

    if os_key == "junos":
        conn.config_mode()
        for ln in lines:
            out = conn.send_command_timing(ln)
            if "error" in out.lower() or "syntax error" in out.lower():
                conn.send_command_timing("rollback 0")
                conn.exit_config_mode()
                raise RuntimeError(f"rejected line '{ln}': {out.strip()[:200]}")
        out = conn.send_command_timing("show | compare")
        if not out.strip():
            conn.send_command_timing("rollback 0")
            conn.exit_config_mode()
            return "junos-nochange"
        res = conn.send_command_timing(f"commit confirmed {confirm_min}")
        if "error" in res.lower():
            conn.send_command_timing("rollback 0")
            conn.exit_config_mode()
            raise RuntimeError(f"commit confirmed failed: {res.strip()[:300]}")
        return "junos-confirm"

    if os_key == "eos":
        session = f"pa-{int(time.time())}"
        conn.send_command_timing(f"configure session {session}")
        for ln in lines:
            out = conn.send_command_timing(ln)
            if "% Invalid" in out or "Error" in out:
                conn.send_command_timing("abort")
                raise RuntimeError(f"rejected line '{ln}': {out.strip()[:200]}")
        out = conn.send_command_timing(f"commit timer {confirm_min:02d}:00:00"[:19])
        if "Error" in out:
            conn.send_command_timing("abort")
            raise RuntimeError(f"commit timer failed: {out.strip()[:300]}")
        return session

    # Platforms without a native confirm: push plainly, restore from snapshot on failure
    out = conn.send_config_set(lines)
    for marker in ("% Invalid", "% Incomplete", "% Ambiguous", "ERROR:", "Invalid input"):
        if marker in out:
            raise RuntimeError(f"device rejected config: {marker} -> {out.strip()[:300]}")
    return "manual"


def confirm(conn, dev: Device, token: str) -> None:
    """Make the change permanent. Must be called before the timer expires."""
    if token == "iosxr-confirm":
        conn.send_command_timing("commit")
    elif token == "junos-confirm":
        conn.config_mode()
        conn.send_command_timing("commit")
        conn.exit_config_mode()
    elif token.startswith("pa-"):           # EOS session
        conn.send_command_timing(f"configure session {token} commit")
    elif token == "manual":
        cmd = SAVE_CMD.get(dev.os)
        if cmd:
            conn.send_command_timing(cmd)


def rollback(conn, dev: Device, token: str, snapshot_id: str | None, disk_cfg: Path) -> None:
    """Return the device to its pre-change state."""
    os_key = dev.os
    LOG.warning("%s: rolling back", dev.host)

    if token == "iosxr-confirm":
        conn.send_command_timing("abort")
        return
    if token == "junos-confirm":
        conn.config_mode()
        conn.send_command_timing("rollback 1")
        conn.send_command_timing("commit")
        conn.exit_config_mode()
        return
    if token.startswith("pa-"):
        conn.send_command_timing(f"configure session {token} abort")
        return

    # Snapshot-based restore
    if os_key == "iosxe" and snapshot_id:
        out = conn.send_command_timing(f"configure replace {snapshot_id} force")
        if "error" in out.lower():
            raise RuntimeError(f"configure replace failed: {out.strip()[:300]}")
        return
    if os_key == "nxos" and snapshot_id:
        out = conn.send_command_timing(f"rollback running-config checkpoint {snapshot_id}")
        if "error" in out.lower() or "fail" in out.lower():
            raise RuntimeError(f"checkpoint rollback failed: {out.strip()[:300]}")
        return
    if os_key == "fortios" and snapshot_id:
        conn.send_command_timing(f"execute restore config flash {snapshot_id}")
        return

    raise RuntimeError(
        f"{dev.host}: no automatic rollback path for {os_key}. "
        f"Restore manually from {disk_cfg}"
    )


def cleanup_snapshot(conn, dev: Device, snapshot_id: str | None) -> None:
    if not snapshot_id:
        return
    try:
        if dev.os == "nxos":
            conn.send_command_timing(f"no checkpoint {snapshot_id}")
        elif dev.os == "iosxe" and snapshot_id.startswith("flash:"):
            conn.send_command_timing(f"delete /force {snapshot_id}")
    except Exception:  # noqa: BLE001 - cleanup must never fail the run
        LOG.debug("%s: snapshot cleanup skipped", dev.host)


# --------------------------------------------------------------------------
# Pre/post assertions
# --------------------------------------------------------------------------

def run_checks(conn, dev: Device) -> dict[str, str]:
    """Execute each declared check and return normalised output keyed by name."""
    results: dict[str, str] = {}
    for chk in dev.checks:
        name, cmd = chk["name"], chk["command"]
        try:
            out = conn.send_command(cmd, read_timeout=60)
            results[name] = " ".join(out.split())
        except Exception as exc:  # noqa: BLE001
            results[name] = f"<error: {exc}>"
    return results


def evaluate(dev: Device, before: dict[str, str], after: dict[str, str]) -> tuple[bool, str]:
    """
    A check passes when either:
      - `expect` is defined and the substring appears in the post-change output, or
      - no `expect` is defined and the output is unchanged from before.
    """
    for chk in dev.checks:
        name = chk["name"]
        exp = chk.get("expect")
        post = after.get(name, "")
        if post.startswith("<error:"):
            return False, f"check '{name}' could not run: {post}"
        if exp:
            if exp not in post:
                return False, f"check '{name}' expected '{exp}', got '{post[:120]}'"
        else:
            if before.get(name) != post:
                return False, f"check '{name}' changed: '{before.get(name, '')[:80]}' -> '{post[:80]}'"
    return True, "all checks passed"


# --------------------------------------------------------------------------
# Per-device workflow
# --------------------------------------------------------------------------

def process(dev: Device, cfg_lines: list[str], args, stamp: str) -> Result:
    res = Result(host=dev.host, os=dev.os)
    t0 = time.time()
    conn = None
    token = ""
    snap_id = None
    disk_path = Path()

    try:
        conn = connect(dev)
        verify_privilege(conn, dev)

        before_cfg = fetch_config(conn, dev)
        disk_path = save_snapshot(before_cfg, dev, stamp, Path(args.snapshot_dir))
        res.snapshot = str(disk_path)

        # Preview the change as a unified diff against the live config
        projected = before_cfg + "\n" + "\n".join(cfg_lines)
        diff = list(difflib.unified_diff(
            before_cfg.splitlines(), projected.splitlines(),
            fromfile="running", tofile="projected", lineterm="", n=2,
        ))
        res.diff_lines = len([d for d in diff if d.startswith(("+", "-")) and not d.startswith(("---", "+++"))])

        if not args.apply:
            res.status = "dry_run"
            res.message = f"{len(cfg_lines)} line(s) would be applied"
            return res

        res.checks_before = run_checks(conn, dev)
        snap_id = onbox_snapshot(conn, dev, f"pa_{stamp}")

        token = push_with_confirm(conn, dev, cfg_lines, args.confirm_minutes)
        if token == "junos-nochange":
            res.status = "skipped"
            res.message = "config already matches intent, nothing to commit"
            return res

        time.sleep(args.settle_seconds)
        res.checks_after = run_checks(conn, dev)
        ok, why = evaluate(dev, res.checks_before, res.checks_after)

        if ok:
            confirm(conn, dev, token)
            cleanup_snapshot(conn, dev, snap_id)
            res.status = "ok"
            res.message = why
        else:
            rollback(conn, dev, token, snap_id, disk_path)
            res.status = "rolled_back"
            res.message = why

    except (NetmikoAuthenticationException,) as exc:
        res.status, res.message = "failed", f"authentication failed: {exc}"
    except (NetmikoTimeoutException, ReadTimeout) as exc:
        res.status = "failed"
        res.message = (
            f"timeout: {exc}. "
            + ("Device will self-revert." if dev.os in NATIVE_CONFIRM
               else f"VERIFY MANUALLY - snapshot at {disk_path}")
        )
    except PermissionError as exc:
        res.status, res.message = "failed", str(exc)
    except Exception as exc:  # noqa: BLE001 - one device must not abort the run
        res.status, res.message = "failed", f"{type(exc).__name__}: {exc}"
        if conn and token:
            try:
                rollback(conn, dev, token, snap_id, disk_path)
                res.status = "rolled_back"
            except Exception as rb:  # noqa: BLE001
                res.message += f" | ROLLBACK ALSO FAILED: {rb}"
    finally:
        res.elapsed_s = round(time.time() - t0, 1)
        if conn:
            try:
                conn.disconnect()
            except Exception:  # noqa: BLE001
                pass
    return res


# --------------------------------------------------------------------------
# Inventory + CLI
# --------------------------------------------------------------------------

def load_inventory(path: str) -> list[Device]:
    raw = Path(path).read_text(encoding="utf-8")
    data = yaml.safe_load(raw) if (yaml and path.endswith((".yaml", ".yml"))) else json.loads(raw)
    devices = data["devices"] if isinstance(data, dict) else data
    return [Device(**d) for d in devices]


def load_config_lines(cfg_dir: str, dev: Device) -> list[str]:
    """Per-host file wins; otherwise the per-OS file; otherwise the common file."""
    base = Path(cfg_dir)
    for candidate in (base / f"{dev.host}.cfg", base / f"{dev.os}.cfg", base / "common.cfg"):
        if candidate.exists():
            return [
                ln.rstrip() for ln in candidate.read_text(encoding="utf-8").splitlines()
                if ln.strip() and not ln.lstrip().startswith("#")
            ]
    return []


def render(results: list[Result]) -> int:
    icon = {"ok": "[OK]", "dry_run": "[DRY]", "rolled_back": "[RBACK]",
            "failed": "[FAIL]", "skipped": "[SKIP]"}
    print("\n" + "=" * 78)
    for r in sorted(results, key=lambda x: x.host):
        print(f"{icon.get(r.status, '[?]'):<8} {r.host:<22} {r.os:<8} "
              f"{r.diff_lines:>3} chg  {r.elapsed_s:>5.1f}s  {r.message}")
    print("=" * 78)

    tally: dict[str, int] = {}
    for r in results:
        tally[r.status] = tally.get(r.status, 0) + 1
    print("  " + "   ".join(f"{k}={v}" for k, v in sorted(tally.items())))

    bad = tally.get("failed", 0) + tally.get("rolled_back", 0)
    if bad:
        print(f"\n  {bad} device(s) need attention. Snapshots kept for every host.")
    return 1 if bad else 0


def main() -> int:
    ap = argparse.ArgumentParser(description="Push config with automatic rollback.")
    ap.add_argument("--inventory", required=True, help="YAML/JSON inventory file")
    ap.add_argument("--config-dir", required=True, help="Directory of .cfg fragments to push")
    ap.add_argument("--apply", action="store_true", help="Actually push (default is dry-run)")
    ap.add_argument("--confirm-minutes", type=int, default=5,
                    help="Auto-revert window on platforms that support it (default 5)")
    ap.add_argument("--settle-seconds", type=int, default=10,
                    help="Wait before post-checks, lets protocols reconverge (default 10)")
    ap.add_argument("--snapshot-dir", default="./snapshots")
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--limit", help="Only hosts whose name contains this substring")
    ap.add_argument("--json-out", help="Write machine-readable results here")
    ap.add_argument("-v", "--verbose", action="store_true")
    args = ap.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)-7s %(message)s",
        datefmt="%H:%M:%S",
    )

    if not os.environ.get("NET_USER") or not os.environ.get("NET_PASS"):
        LOG.error("Set NET_USER and NET_PASS in the environment.")
        return 2

    devices = load_inventory(args.inventory)
    if args.limit:
        devices = [d for d in devices if args.limit in d.host]
    if not devices:
        LOG.error("No devices matched.")
        return 2

    stamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
    mode = "APPLY" if args.apply else "DRY-RUN"
    LOG.info("%s | %d device(s) | confirm window %d min | snapshot %s",
             mode, len(devices), args.confirm_minutes, stamp)

    if args.apply:
        LOG.warning("Live change. Devices on %s self-revert if this script dies.",
                    "/".join(sorted(NATIVE_CONFIRM)))

    work: dict[Any, Device] = {}
    results: list[Result] = []
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        for dev in devices:
            lines = load_config_lines(args.config_dir, dev)
            if not lines:
                results.append(Result(host=dev.host, os=dev.os, status="skipped",
                                      message="no config fragment found"))
                continue
            work[pool.submit(process, dev, lines, args, stamp)] = dev
        for fut in as_completed(work):
            dev = work[fut]
            try:
                results.append(fut.result())
            except Exception as exc:  # noqa: BLE001
                LOG.exception("%s: worker crashed", dev.host)
                results.append(Result(host=dev.host, os=dev.os, status="failed",
                                      message=f"worker crashed: {exc}"))

    if args.json_out:
        Path(args.json_out).write_text(
            json.dumps([asdict(r) for r in results], indent=2), encoding="utf-8")
        LOG.info("Results written to %s", args.json_out)

    return render(results)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\nInterrupted. Devices inside a confirm window will self-revert.", file=sys.stderr)
        sys.exit(130)