#!/usr/bin/env python3
"""Peer the CML core routers with the GoBGP full-table feed (roadmap E1).

GoBGP (AS65001, 10.40.40.250) holds the full real Internet table pulled from
the Bromirski route server. This script configures CORE-01/CORE-02 (AS65020)
to peer eBGP with GoBGP and accept that table. As route reflectors the cores
then propagate it to every R9K client -- so all 9 lab routers carry and
BMP-export a full table. This is an intentional lab stress test of the
OpenBMP ingestion/storage path.

Applied per core (additive -- no existing session/policy is modified):
  * route-policy GOBGP-FEED-PASS  (a plain `pass` policy; eBGP needs one)
  * neighbor 10.40.40.202 remote-as 65001, ebgp-multihop, mgmt-sourced,
    IPv4-unicast only, with a maximum-prefix safety cap.

The matching GoBGP side is gobgp/gobgpd.conf (neighbors 10.100.0.100/.200);
restart GoBGP after applying:  docker compose up -d gobgp

IOS-XR BMP config is not exposed via NETCONF on this release, so -- like
cml/proxmox_bmp_config.py -- this applies config over the SSH CLI.

Covers both labs: CML cores (AS65020) and PROX cores (AS65021).

Usage:
    python3 cml/gobgp_peering_config.py                 # apply, all 4 cores
    python3 cml/gobgp_peering_config.py cml             # apply, CML cores only
    python3 cml/gobgp_peering_config.py prox            # apply, PROX cores only
    python3 cml/gobgp_peering_config.py --remove        # ROLLBACK, all cores
    python3 cml/gobgp_peering_config.py --remove prox   # ROLLBACK, PROX only

Rollback: `--remove` deletes the GoBGP neighbor and the GOBGP-FEED-PASS
policy from the cores. To stop the feed instantly without touching router
config, `docker compose stop gobgp` -- the eBGP sessions drop and the full
table is withdrawn fleet-wide within seconds. See gobgp/README.md.
"""
import os
import sys
import time
import paramiko


def _env_default(key, default, dotenv=".env"):
    """Resolve a value from os.environ or the repo-root .env, else default."""
    v = os.environ.get(key)
    if v:
        return v
    try:
        with open(dotenv) as fh:
            for line in fh:
                s = line.strip()
                if s and not s.startswith("#") and s.startswith(f"{key}="):
                    return s.split("=", 1)[1].strip().strip('"').strip("'")
    except FileNotFoundError:
        pass
    return default


# GoBGP runs network_mode: host, so it sources BGP TCP from the host's real
# interface IP -- NOT its router-id. The cores must peer with the host IP.
# Resolved from $HOST_IP or the HOST_IP= line in repo-root .env.
GOBGP_IP = _env_default("HOST_IP", "10.40.40.202")
GOBGP_AS = "65001"

# Additive config, built per core (asn = that core's local BGP AS:
# CML lab = 65020, PROX lab = 65021). Flat formal-form lines applied at the
# (config)# prompt.
# IPv4-unicast only: the cores have no global IPv6 address, so an ipv6-unicast
# AF on this IPv4-transport session holds the whole neighbor Idle. The IPv6
# full-table feed is a separate phase (needs a v6-transport session or v6
# addressing on the cores).
def apply_lines(asn):
    n = f"router bgp {asn} neighbor {GOBGP_IP}"
    return [
        "route-policy GOBGP-FEED-PASS",
        "  pass",
        "end-policy",
        f"{n} remote-as {GOBGP_AS}",
        f"{n} description GoBGP full-table feed (lab stress test)",
        f"{n} ebgp-multihop 64",
        f"{n} update-source MgmtEth0/RP0/CPU0/0",
        f"{n} address-family ipv4 unicast route-policy GOBGP-FEED-PASS in",
        f"{n} address-family ipv4 unicast route-policy GOBGP-FEED-PASS out",
        f"{n} address-family ipv4 unicast maximum-prefix 1500000 90",
    ]


# Rollback -- remove the neighbor (and its sub-config) then the policy.
def remove_lines(asn):
    return [
        f"no router bgp {asn} neighbor {GOBGP_IP}",
        "no route-policy GOBGP-FEED-PASS",
    ]


# (name, mgmt_ip, user, password, local_asn) -- both labs.
CORES = [
    ("CML-CORE-01",  "10.100.0.100", "webui", "cisco", "65020"),
    ("CML-CORE-02",  "10.100.0.200", "webui", "cisco", "65020"),
    ("PROX-CORE-01", "10.100.1.100", "admin", "cisco", "65021"),
    ("PROX-CORE-02", "10.100.1.200", "admin", "cisco", "65021"),
]


def _drain(shell, settle=0.4, limit=20.0, until=None):
    out, start = "", time.time()
    while time.time() - start < limit:
        time.sleep(settle)
        if shell.recv_ready():
            out += shell.recv(65535).decode(errors="replace")
            if until and until in out:
                break
        elif until is None:
            break
        elif until in out:
            break
    return out


def configure_core(name, ip, user, pwd, asn, mode):
    verb = "applying" if mode == "apply" else "removing"
    lines = apply_lines(asn) if mode == "apply" else remove_lines(asn)
    print(f"\n=== {name} ({ip}) AS{asn} -- {verb} GoBGP peering ===")
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(ip, username=user, password=pwd, timeout=15,
                    look_for_keys=False, allow_agent=False)
        shell = ssh.invoke_shell(width=220, height=2000)
        time.sleep(2)
        shell.recv(65535)
        CFG = "(config)#"

        shell.send("terminal length 0\n")
        _drain(shell, 0.4, 5)
        shell.send("configure terminal\n")
        out = _drain(shell, 0.4, 15, until=CFG)
        if CFG not in out:
            print(f"  FAIL: could not enter config mode\n  {out[-200:]}")
            ssh.close()
            return False

        for line in lines:
            shell.send(line + "\n")
            time.sleep(0.4)
        _drain(shell, 0.3, 8, until=CFG)

        shell.send("show configuration\n")
        cand = _drain(shell, 0.3, 10, until=CFG)
        if GOBGP_IP not in cand and "GOBGP-FEED-PASS" not in cand:
            print(f"  OK: no changes ({mode} already in effect)")
            shell.send("abort\n")
            _drain(shell, 0.5, 5)
            ssh.close()
            return True

        shell.send("commit\n")
        result = _drain(shell, 0.3, 25, until=CFG)
        if "fail" in result.lower() or "error" in result.lower():
            print(f"  FAIL: commit error\n  {result[-300:]}")
            shell.send("abort\n")
            _drain(shell, 0.5, 5)
            ssh.close()
            return False

        shell.send("end\n")
        _drain(shell, 1.0, 8)
        if mode == "apply":
            shell.send(f"show bgp ipv4 unicast neighbors {GOBGP_IP} | include BGP state\n")
            verify = _drain(shell, 1.0, 12)
            state = next((l.strip() for l in verify.splitlines()
                          if "BGP state" in l), "(state not yet reported)")
            print(f"  committed. {state}")
        else:
            shell.send(f"show running-config router bgp | include {GOBGP_IP}\n")
            verify = _drain(shell, 1.0, 12)
            gone = GOBGP_IP not in verify.replace(f"include {GOBGP_IP}", "")
            print(f"  committed. neighbor removed: {gone}")
        ssh.close()
        return True
    except Exception as e:
        print(f"  FAIL: {e}")
        return False


def main():
    args = [a for a in sys.argv[1:]]
    mode = "apply"
    if "--remove" in args:
        mode = "remove"
        args.remove("--remove")
    target = args[0].lower() if args else None

    if mode == "remove":
        print("ROLLBACK: removing GoBGP peering from the core routers.")
    results = {}
    for name, ip, user, pwd, asn in CORES:
        if target and target not in name.lower():
            continue
        results[name] = configure_core(name, ip, user, pwd, asn, mode)

    print(f"\n{'='*48}\n  SUMMARY ({mode})")
    for name, ok in results.items():
        print(f"  {name:22s} {'OK' if ok else 'FAILED'}")
    if mode == "apply":
        print("\nNext: restart GoBGP to load the new neighbors:")
        print("  docker compose up -d gobgp")
    else:
        print("\nGoBGP container config still lists the cores; that is inert")
        print("with the neighbors removed. To fully revert, also restore the")
        print("previous gobgp/gobgpd.conf and run: docker compose up -d gobgp")
    sys.exit(0 if all(results.values()) else 1)


if __name__ == "__main__":
    main()