BGP Prefix Count Monitor & Leak Detector

#!/usr/bin/env python3
"""
BGP Prefix Count Monitor & Leak Detector

READ THE DOCUMENTATION BEFORE YOU RUN THIS:
    https://permitany.com/scripts/bgp-prefix-monitor

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 csv
import re
import time
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
# ---------------------------------------------------------------------------

BGP_CMDS = {
    "ios":     {"global": "show ip bgp summary",
                "vrf": "show ip bgp vpnv4 vrf {vrf} summary"},
    "iosxe":   {"global": "show ip bgp summary",
                "vrf": "show ip bgp vpnv4 vrf {vrf} summary"},
    "iosxr":   {"global": "show bgp summary",
                "vrf": "show bgp vrf {vrf} summary"},
    "nxos":    {"global": "show bgp summary",
                "vrf": "show bgp vrf {vrf} summary"},
    "junos":   {"global": "show bgp summary",
                "vrf": "show bgp summary instance {vrf}"},
    "eos":     {"global": "show ip bgp summary",
                "vrf": "show ip bgp vrf {vrf} summary"},
    "panos":   {"global": "show routing protocol bgp peer",
                "vrf": "show routing protocol bgp peer virtual-router {vrf}"},
    "fortios": {"global": "get router info bgp summary",
                "vrf": "get router info bgp summary vrf {vrf}"},
}

IPV4 = re.compile(r"^\d{1,3}(?:\.\d{1,3}){3}$")
NOT_ESTABLISHED = re.compile(
    r"\b(Idle|Active|Connect|OpenSent|OpenConfirm|Established)\b", re.IGNORECASE)


# ---------------------------------------------------------------------------
# Parsing
# ---------------------------------------------------------------------------

def parse_prefix_counts(output: str, os_key: str) -> dict[str, dict]:
    """
    Return {peer_ip: {"prefixes": int, "state": str}}.

    The last column of a BGP summary is State/PfxRcd: a number when the session
    is Established, a state word otherwise. Treating a state word as zero
    prefixes would look like a total withdrawal, so the two are distinguished.
    """
    peers: dict[str, dict] = {}

    if os_key == "junos":
        # Peer  AS  InPkt  OutPkt  OutQ  Flaps  Last Up/Dwn  State|#Active/Received/Accepted/Damped
        for line in output.splitlines():
            parts = line.split()
            if len(parts) < 7 or not IPV4.match(parts[0]):
                continue
            tail = parts[-1]
            m = re.match(r"(\d+)/(\d+)/(\d+)/(\d+)", tail)
            if m:
                peers[parts[0]] = {"prefixes": int(m.group(2)), "state": "Established"}
            else:
                peers[parts[0]] = {"prefixes": None, "state": tail}
        return peers

    if os_key == "panos":
        for line in output.splitlines():
            parts = line.split()
            if len(parts) < 3 or not IPV4.match(parts[0]):
                continue
            state = next((p for p in parts if NOT_ESTABLISHED.fullmatch(p)), "unknown")
            nums = [int(p) for p in parts if p.isdigit()]
            peers[parts[0]] = {"prefixes": nums[-1] if nums else None, "state": state}
        return peers

    # IOS / IOS-XE / IOS-XR / NX-OS / EOS / FortiOS share the summary layout.
    for line in output.splitlines():
        parts = line.split()
        if len(parts) < 8 or not IPV4.match(parts[0]):
            continue
        last = parts[-1]
        if last.isdigit():
            peers[parts[0]] = {"prefixes": int(last), "state": "Established"}
        else:
            peers[parts[0]] = {"prefixes": None, "state": last}
    return peers


# ---------------------------------------------------------------------------
# Baseline
# ---------------------------------------------------------------------------
# The baseline is the rolling average of recent samples. The rule that makes
# it useful: a sample that triggered an alert is NOT folded into the baseline.
# Otherwise a route leak is absorbed within a few polls and becomes the new
# "normal", which is exactly when you stop being told about it.

def load_baselines(path: str) -> dict:
    try:
        return json.loads(Path(path).read_text(encoding="utf-8"))
    except FileNotFoundError:
        return {}
    except json.JSONDecodeError as exc:
        raise SystemExit(f"Baseline file {path} is corrupt: {exc}\n"
                         f"Delete it to start learning again from scratch.")


def save_baselines(baselines: dict, path: str) -> None:
    tmp = Path(path).with_suffix(".tmp")
    tmp.write_text(json.dumps(baselines, indent=2), encoding="utf-8")
    tmp.replace(path)          # atomic, so a crash cannot truncate the file


def evaluate(key: str, count: int, baselines: dict, dev: dict, args) -> tuple[float, list[str]]:
    """Compare this sample against the baseline. Returns (baseline, alerts)."""
    entry = baselines.get(key, {"samples": [], "avg": 0.0, "frozen": False})
    samples = entry.get("samples", [])
    avg = entry.get("avg", 0.0)
    alerts: list[str] = []

    limits = dev.get("thresholds", {})
    lo, hi = limits.get("min"), limits.get("max")
    if lo is not None and count < lo:
        alerts.append(f"BELOW_MIN expected >={lo:,}, got {count:,}")
    if hi is not None and count > hi:
        alerts.append(f"ABOVE_MAX expected <={hi:,}, got {count:,}")

    if len(samples) < args.learn_samples:
        return avg, alerts          # still learning; no baseline comparison yet

    if avg > 0:
        if count > avg * args.leak_multiplier:
            alerts.append(f"LEAK {count:,} prefixes, {(count/avg-1)*100:.0f}% above "
                          f"baseline {avg:,.0f}")
        elif count < avg * (1 - args.drop_pct):
            alerts.append(f"DROP {count:,} prefixes, {(avg-count)/avg*100:.1f}% below "
                          f"baseline {avg:,.0f}")
    return avg, alerts


def update_baseline(key: str, count: int, baselines: dict, alerted: bool, args) -> None:
    """Fold a sample into the baseline, unless it was anomalous."""
    entry = baselines.setdefault(key, {"samples": [], "avg": 0.0, "frozen": False})
    if alerted:
        entry["frozen"] = True
        entry["last_anomaly"] = count
        return                      # an anomaly must never become the new normal
    entry["frozen"] = False
    entry["samples"] = (entry["samples"] + [count])[-args.baseline_window:]
    entry["avg"] = sum(entry["samples"]) / len(entry["samples"])


# ---------------------------------------------------------------------------
# Polling
# ---------------------------------------------------------------------------

def poll_device(dev: dict, args) -> list[dict]:
    """Collect raw samples. Does not touch the shared baseline dict - that is
    done single-threaded after every worker has finished, to avoid a race."""
    host = dev["host"]
    os_key = dev["os"]
    cmds = BGP_CMDS[os_key]
    conn = None
    rows: list[dict] = []
    ts = datetime.now(timezone.utc).isoformat(timespec="seconds")

    try:
        conn = connect(dev)
        for vrf in [None] + list(dev.get("vrfs", [])):
            scope = vrf or "global"
            cmd = cmds["global"] if vrf is None else cmds["vrf"].format(vrf=vrf)
            out = conn.send_command(cmd, read_timeout=args.read_timeout)
            for peer, info in parse_prefix_counts(out, os_key).items():
                rows.append({"ts": ts, "host": host, "os": os_key, "vrf": scope,
                             "peer": peer, "prefixes": info["prefixes"],
                             "state": info["state"], "thresholds": dev.get("thresholds", {})})
        return rows

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


def append_csv(rows: list[dict], path: str) -> None:
    new = not Path(path).exists()
    with open(path, "a", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=["ts", "host", "vrf", "peer", "state",
                                          "prefixes", "baseline", "alerts"])
        if new:
            w.writeheader()
        for r in rows:
            w.writerow({
                "ts": r["ts"], "host": r["host"], "vrf": r["vrf"], "peer": r["peer"],
                "state": r.get("state", ""),
                "prefixes": "" if r.get("prefixes") is None else r["prefixes"],
                "baseline": f"{r.get('baseline', 0):.0f}",
                "alerts": "|".join(r.get("alerts", [])),
            })


def run_once(devices: list[dict], baselines: dict, args) -> int:
    print(f"\nPolling {len(devices)} device(s) at "
          f"{datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')} UTC")
    rows: list[dict] = []
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(poll_device, d, args): d for d in devices}
        for fut in as_completed(futures):
            dev = futures[fut]
            try:
                rows.extend(fut.result())
            except Exception as exc:  # noqa: BLE001
                rows.append({"ts": "", "host": dev["host"], "os": dev.get("os", "?"),
                             "vrf": "-", "peer": "-", "prefixes": None,
                             "state": "ERROR", "error": f"worker crashed: {exc}"})

    # Single-threaded: evaluating and updating the shared baseline dict from
    # inside the workers would be a data race.
    dev_by_host = {d["host"]: d for d in devices}
    for r in rows:
        if r.get("state") == "ERROR" or r["prefixes"] is None:
            r["baseline"], r["alerts"] = 0.0, []
            continue
        key = f"{r['host']}|{r['vrf']}|{r['peer']}"
        base, alerts = evaluate(key, r["prefixes"], baselines,
                                dev_by_host.get(r["host"], {}), args)
        update_baseline(key, r["prefixes"], baselines, bool(alerts), args)
        r["baseline"], r["alerts"] = base, alerts

    for r in sorted(rows, key=lambda x: (x["host"], x["vrf"], x["peer"])):
        scope = "[global]" if r["vrf"] == "global" else f"[vrf:{r['vrf']}]"
        if r.get("state") == "ERROR":
            print(f"[ERROR ] {r['host']:<20} {r.get('error', '')}")
            continue
        if r["prefixes"] is None:
            print(f"[DOWN  ] {r['host']:<20} {scope:<14} peer {r['peer']:<16} "
                  f"session is {r['state']}, not Established")
            continue
        head = f"{r['host']:<20} {scope:<14} peer {r['peer']:<16} "
        body = f"{r['prefixes']:>9,} prefixes  (baseline {r['baseline']:>9,.0f})"
        if r["alerts"]:
            for a in r["alerts"]:
                print(f"[{a.split()[0]:<6}] {head}{body}")
                print(f"           {a}")
        else:
            print(f"[OK    ] {head}{body}")

    save_baselines(baselines, args.baseline_file)
    if args.csv:
        append_csv(rows, args.csv)

    leaks = sum(1 for r in rows if any(a.startswith("LEAK") for a in r.get("alerts", [])))
    drops = sum(1 for r in rows if any(a.startswith("DROP") for a in r.get("alerts", [])))
    thresh = sum(1 for r in rows
                 if any(a.startswith(("BELOW", "ABOVE")) for a in r.get("alerts", [])))
    down = sum(1 for r in rows if r["prefixes"] is None and r.get("state") != "ERROR")
    errs = sum(1 for r in rows if r.get("state") == "ERROR")
    learning = sum(1 for r in rows
                   if len(baselines.get(f"{r['host']}|{r['vrf']}|{r['peer']}", {})
                          .get("samples", [])) < args.learn_samples)

    print(f"\nLeaks {leaks}  Drops {drops}  Threshold {thresh}  "
          f"Sessions down {down}  Unreachable {errs}")
    if learning:
        print(f"{learning} peer(s) still learning - no baseline comparison until "
              f"{args.learn_samples} clean samples have been seen.")
    frozen = [k for k, v in baselines.items() if v.get("frozen")]
    if frozen:
        print(f"{len(frozen)} baseline(s) frozen because the last sample was "
              f"anomalous - they will not absorb the anomaly.")

    return 1 if (leaks or drops or thresh or errs) else 0


def main() -> int:
    ap = argparse.ArgumentParser(
        description="Monitor BGP prefix counts per peer and detect leaks and withdrawals.",
        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 poll")
    ap.add_argument("--once", action="store_true",
                    help="Poll once and exit. Without it the script runs continuously")
    ap.add_argument("--interval", type=int, default=300,
                    help="Seconds between polls in continuous mode. Default: 300")
    ap.add_argument("--leak-multiplier", type=float, default=3.0,
                    help="Alert when the count exceeds baseline x this. Default: 3.0")
    ap.add_argument("--drop-pct", type=float, default=0.15,
                    help="Alert when the count falls this fraction below baseline. "
                         "Default: 0.15 (15 percent)")
    ap.add_argument("--baseline-window", type=int, default=12,
                    help="How many clean samples the rolling average covers. Default: 12")
    ap.add_argument("--learn-samples", type=int, default=3,
                    help="Clean samples required before a baseline is trusted. "
                         "Default: 3")
    ap.add_argument("--baseline-file", default="bgp_baselines.json",
                    help="Where the learned baselines persist between runs")
    ap.add_argument("--csv", help="Append every sample to this CSV, for trending")
    ap.add_argument("--limit", help="Only poll hosts whose name contains this text")
    ap.add_argument("--read-timeout", type=int, default=90,
                    help="Seconds to wait per command. Default: 90")
    ap.add_argument("--workers", type=int, default=10,
                    help="How many devices to poll at once. Default: 10")
    args = ap.parse_args()

    devices = load_inventory(args.inventory)
    if args.limit:
        devices = [d for d in devices if args.limit in d["host"]]
        if not devices:
            raise SystemExit(f"No device name contains '{args.limit}'.")

    credentials()  # fail fast before opening any connection

    baselines = load_baselines(args.baseline_file)
    print(f"Loaded {len(baselines)} baseline(s) from {args.baseline_file}")

    if args.once:
        return run_once(devices, baselines, args)

    while True:
        run_once(devices, baselines, args)
        print(f"Next poll in {args.interval}s. Ctrl-C to stop.")
        time.sleep(args.interval)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\nStopped. Baselines were saved after the last completed poll.",
              file=sys.stderr)
        sys.exit(130)