BGP Community Tagger

#!/usr/bin/env python3
"""
BGP Community Tagger

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

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 ipaddress
import re

# ---------------------------------------------------------------------------
# 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


# ---------------------------------------------------------------------------
# Community handling
# ---------------------------------------------------------------------------

WELL_KNOWN = {
    "no-export": "no-export",
    "no-advertise": "no-advertise",
    "local-as": "local-AS",
    "no-export-subconfed": "local-AS",
    "graceful-shutdown": "graceful-shutdown",
    "blackhole": "blackhole",
}

COMMUNITY_RE = re.compile(r"^(\d{1,5}:\d{1,5}|[\d]{1,10}|no-export|no-advertise|"
                          r"local-AS|graceful-shutdown|blackhole)$", re.IGNORECASE)


def normalise_community(value: str) -> str:
    return WELL_KNOWN.get(str(value).lower(), str(value))


def validate_rules(rules: list[dict]) -> None:
    """Catch malformed rules before any config is generated or pushed."""
    if not rules:
        raise SystemExit("No tagging rules found. The file needs a `rules:` list.")
    for i, r in enumerate(rules, 1):
        if not isinstance(r, dict):
            raise SystemExit(f"Rule #{i} is not a mapping.")
        for key in ("prefix", "community", "action"):
            if key not in r:
                raise SystemExit(f"Rule #{i} is missing `{key}:`")
        try:
            ipaddress.ip_network(r["prefix"], strict=False)
        except ValueError as exc:
            raise SystemExit(f"Rule #{i}: invalid prefix '{r['prefix']}' - {exc}")
        if r["action"] not in ("add", "remove", "replace"):
            raise SystemExit(f"Rule #{i}: action must be add, remove or replace, "
                             f"got '{r['action']}'")
        if not COMMUNITY_RE.match(str(r["community"])):
            raise SystemExit(
                f"Rule #{i}: '{r['community']}' is not a valid community.\n"
                f"Use ASN:VALUE (e.g. 65000:100) or a well-known name "
                f"(no-export, no-advertise, local-as, graceful-shutdown).")


# ---------------------------------------------------------------------------
# Config generation, per platform
# ---------------------------------------------------------------------------

def generate_ios(rules: list[dict], rm: str, pl: str) -> list[str]:
    """
    IOS / IOS-XE / NX-OS / EOS route-map.

    Each rule gets its OWN prefix-list and its own route-map sequence. The
    original approach of one shared prefix-list made every sequence match
    every prefix, so whichever community came first won for all of them.
    """
    out: list[str] = []
    for i, r in enumerate(rules, 1):
        name = f"{pl}-{i}"
        out.append(f"ip prefix-list {name} seq 10 permit {r['prefix']}")
    out.append("!")

    for i, r in enumerate(rules, 1):
        seq = i * 10
        comm = normalise_community(r["community"])
        out.append(f"route-map {rm} permit {seq}")
        out.append(f" match ip address prefix-list {pl}-{i}")
        if r["action"] == "add":
            out.append(f" set community {comm} additive")
        elif r["action"] == "remove":
            out.append(f" set comm-list {rm}-DEL delete")
        else:
            out.append(f" set community {comm}")
        out.append("!")

    # A route-map ends with an implicit deny; without this every prefix not
    # matched above would be dropped by the policy.
    out.append(f"route-map {rm} permit 65535")
    out.append("!")

    removals = [normalise_community(r["community"]) for r in rules
                if r["action"] == "remove"]
    if removals:
        out.insert(0, f"ip community-list expanded {rm}-DEL permit "
                      f"{'|'.join(removals)}")
        out.insert(1, "!")
    return out


def generate_iosxr(rules: list[dict], policy: str) -> list[str]:
    out = [f"route-policy {policy}"]
    for i, r in enumerate(rules):
        comm = normalise_community(r["community"])
        kw = "if" if i == 0 else "elseif"
        out.append(f"  {kw} destination in ({r['prefix']}) then")
        if r["action"] == "add":
            out.append(f"    set community ({comm}) additive")
        elif r["action"] == "remove":
            out.append(f"    delete community in ({comm})")
        else:
            out.append(f"    set community ({comm})")
    out.append("  endif")
    out.append("  pass")        # without this, unmatched routes are dropped
    out.append("end-policy")
    return out


def generate_junos(rules: list[dict], policy: str) -> list[str]:
    out: list[str] = []
    seen: set[str] = set()
    for r in rules:
        comm = normalise_community(r["community"])
        cname = "COMM-" + re.sub(r"[^A-Za-z0-9]", "-", comm).upper()
        if cname not in seen:
            out.append(f"set policy-options community {cname} members {comm}")
            seen.add(cname)

    for i, r in enumerate(rules, 1):
        comm = normalise_community(r["community"])
        cname = "COMM-" + re.sub(r"[^A-Za-z0-9]", "-", comm).upper()
        term = f"TERM-{i}"
        out.append(f"set policy-options policy-statement {policy} term {term} "
                   f"from route-filter {r['prefix']} exact")
        verb = {"add": "add", "remove": "delete", "replace": "set"}[r["action"]]
        out.append(f"set policy-options policy-statement {policy} term {term} "
                   f"then community {verb} {cname}")
        out.append(f"set policy-options policy-statement {policy} term {term} "
                   f"then accept")

    out.append(f"set policy-options policy-statement {policy} term DEFAULT then accept")
    return out


def generate(os_key: str, rules: list[dict], args) -> tuple[list[str], str]:
    if os_key in ("ios", "iosxe"):
        return generate_ios(rules, args.route_map, args.prefix_list), "IOS/IOS-XE route-map"
    if os_key == "nxos":
        return generate_ios(rules, args.route_map, args.prefix_list), "NX-OS route-map"
    if os_key == "eos":
        return generate_ios(rules, args.route_map, args.prefix_list), "Arista EOS route-map"
    if os_key == "iosxr":
        return generate_iosxr(rules, args.policy), "IOS-XR route-policy"
    if os_key == "junos":
        return generate_junos(rules, args.policy), "Junos policy-statement"
    raise ValueError(f"BGP community tagging is not generated for os '{os_key}'")


UNSUPPORTED = {
    "panos": "PAN-OS BGP policy is configured through its own object model",
    "fortios": "FortiOS route-map syntax differs enough to need its own template",
}


# ---------------------------------------------------------------------------
# Push with verification
# ---------------------------------------------------------------------------

VERIFY_CMD = {
    "ios": "show route-map {rm}",
    "iosxe": "show route-map {rm}",
    "nxos": "show route-map {rm}",
    "eos": "show route-map {rm}",
    "iosxr": "show rpl route-policy {policy}",
    "junos": "show configuration policy-options policy-statement {policy}",
}


def apply_to_device(dev: dict, lines: list[str], args) -> dict:
    host = dev["host"]
    os_key = dev["os"]
    conn = None
    try:
        conn = connect(dev)

        if os_key == "iosxr":
            conn.send_config_set(lines, exit_config_mode=False)
            out = conn.send_command_timing(f"commit confirmed {args.confirm_minutes} minutes")
            if "rror" in out or "ailed" in out:
                conn.send_command_timing("abort")
                return {"host": host, "os": os_key, "status": "failed",
                        "error": f"commit confirmed rejected: {out.strip()[:200]}"}
            conn.exit_config_mode()
            token = "iosxr"
        elif os_key == "junos":
            conn.config_mode()
            for ln in lines:
                r = conn.send_command_timing(ln)
                if "error" in r.lower():
                    conn.send_command_timing("rollback 0")
                    conn.exit_config_mode()
                    return {"host": host, "os": os_key, "status": "failed",
                            "error": f"rejected '{ln}': {r.strip()[:160]}"}
            r = conn.send_command_timing(f"commit confirmed {args.confirm_minutes}")
            if "error" in r.lower():
                conn.send_command_timing("rollback 0")
                conn.exit_config_mode()
                return {"host": host, "os": os_key, "status": "failed",
                        "error": f"commit confirmed failed: {r.strip()[:200]}"}
            token = "junos"
        else:
            out = conn.send_config_set(lines, read_timeout=120)
            for marker in ("% Invalid", "% Incomplete", "% Ambiguous", "ERROR:"):
                if marker in out:
                    return {"host": host, "os": os_key, "status": "failed",
                            "error": f"device rejected config ({marker})"}
            token = "plain"

        # Verify the policy exists before confirming anything
        vcmd = VERIFY_CMD[os_key].format(rm=args.route_map, policy=args.policy)
        verify = conn.send_command(vcmd, read_timeout=60)
        target = args.policy if os_key in ("iosxr", "junos") else args.route_map
        if target not in verify:
            if token == "iosxr":
                conn.send_command_timing("abort")
            elif token == "junos":
                conn.config_mode()
                conn.send_command_timing("rollback 1")
                conn.send_command_timing("commit")
                conn.exit_config_mode()
            return {"host": host, "os": os_key, "status": "rolled_back",
                    "error": f"'{target}' not present after push - change reverted"}

        if token == "iosxr":
            conn.send_command_timing("commit")
        elif token == "junos":
            conn.config_mode()
            conn.send_command_timing("commit")
            conn.exit_config_mode()
        elif os_key in ("ios", "iosxe", "eos"):
            conn.send_command_timing("write memory")
        elif os_key == "nxos":
            conn.send_command_timing("copy running-config startup-config")

        return {"host": host, "os": os_key, "status": "ok",
                "lines": len(lines), "verified": target}

    except NetmikoAuthenticationException as exc:
        return {"host": host, "os": os_key, "status": "failed",
                "error": f"authentication failed - check NET_USER/NET_PASS ({exc})"}
    except NetmikoTimeoutException as exc:
        note = ("device will self-revert" if os_key in ("iosxr", "junos")
                else "VERIFY DEVICE STATE MANUALLY")
        return {"host": host, "os": os_key, "status": "failed",
                "error": f"timeout: {exc} - {note}"}
    except Exception as exc:  # noqa: BLE001 - one device must not abort the run
        return {"host": host, "os": os_key, "status": "failed",
                "error": f"{type(exc).__name__}: {exc}"}
    finally:
        if conn:
            try:
                conn.disconnect()
            except Exception:  # noqa: BLE001
                pass


def load_rules(path: str) -> list[dict]:
    p = Path(path)
    if not p.exists():
        raise SystemExit(
            f"Rules file not found: {path}\n"
            f"Create it with:\n\n"
            f"  rules:\n"
            f"    - prefix: 10.0.0.0/8\n"
            f"      community: \"65000:100\"\n"
            f"      action: add\n")
    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}")
    rules = data.get("rules", data) if isinstance(data, dict) else data
    validate_rules(rules)
    return rules


def main() -> int:
    ap = argparse.ArgumentParser(
        description="Generate and optionally apply BGP community tagging policy.",
        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 routers to target")
    ap.add_argument("--rules", required=True,
                    help="YAML/JSON file of prefix/community/action tagging rules")
    ap.add_argument("--apply", action="store_true",
                    help="Push the generated config. Omit for dry-run, the default")
    ap.add_argument("--confirm-minutes", type=int, default=5,
                    help="Auto-revert window on IOS-XR and Junos. Default: 5")
    ap.add_argument("--route-map", default="SET-BGP-COMMUNITIES",
                    help="Route-map name on IOS/IOS-XE/NX-OS/EOS")
    ap.add_argument("--prefix-list", default="PFX-COMMUNITY-TAG",
                    help="Prefix-list name prefix; each rule gets its own list")
    ap.add_argument("--policy", default="BGP-COMMUNITY-POLICY",
                    help="Policy name on IOS-XR and Junos")
    ap.add_argument("--limit", help="Only target hosts whose name contains this text")
    ap.add_argument("--out-dir", help="Write the generated config per device to files here")
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--json-out", help="Write machine-readable results here")
    args = ap.parse_args()

    rules = load_rules(args.rules)
    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}'.")

    if args.apply:
        credentials()

    mode = "APPLY" if args.apply else "DRY-RUN"
    print(f"{mode} | {len(rules)} tagging rule(s) | {len(devices)} device(s)")
    if args.apply:
        print("Live change. IOS-XR and Junos devices self-revert if this script dies.")
    print("=" * 80)

    generated: list[tuple[dict, list[str], str]] = []
    for dev in devices:
        try:
            lines, kind = generate(dev["os"], rules, args)
        except ValueError:
            note = UNSUPPORTED.get(dev["os"], f"not generated for {dev['os']}")
            print(f"[SKIP]  {dev['host']} ({dev['os']})  {note}")
            continue
        generated.append((dev, lines, kind))

        print(f"\n{dev['host']} ({dev['os']})  -  {kind}, {len(lines)} line(s)")
        if not args.apply:
            for ln in lines:
                print(f"    {ln}")

        if args.out_dir:
            d = Path(args.out_dir)
            d.mkdir(parents=True, exist_ok=True)
            (d / f"{dev['host']}.cfg").write_text("\n".join(lines) + "\n", encoding="utf-8")

    if args.out_dir and generated:
        print(f"\nWrote {len(generated)} config file(s) to {args.out_dir}/")
        print("Review them, then lint with config_validator.py before pushing.")

    if not args.apply:
        print("\n" + "=" * 80)
        print("DRY-RUN. Nothing was sent. Re-run with --apply to push.")
        return 0

    results = []
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(apply_to_device, dev, lines, args): dev["host"]
                   for dev, lines, _ in generated}
        for fut in as_completed(futures):
            host = futures[fut]
            try:
                results.append(fut.result())
            except Exception as exc:  # noqa: BLE001
                results.append({"host": host, "os": "?", "status": "failed",
                                "error": f"worker crashed: {exc}"})

    print("\n" + "=" * 80)
    for r in sorted(results, key=lambda x: x["host"]):
        tag = {"ok": "[OK]", "failed": "[FAIL]", "rolled_back": "[RBACK]"}[r["status"]]
        detail = (f"{r['lines']} line(s) applied, '{r['verified']}' verified on device"
                  if r["status"] == "ok" else r.get("error", ""))
        print(f"{tag:<8} {r['host']:<24} {r['os']:<8} {detail}")

    ok = sum(1 for r in results if r["status"] == "ok")
    bad = len(results) - ok
    print(f"\n{ok} applied and verified, {bad} failed or reverted.")

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

    return 1 if bad else 0


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)