netbox-diode-project/collectors/network_collector.py

#!/usr/bin/env python3
"""Network device collector for NetBox via Diode SDK.

Discovers Cisco and Brocade network devices via NAPALM (and optionally pyATS
for CDP/OSPF/IS-IS), then ingests devices, interfaces, IPs, cables, VLANs,
VRFs, prefixes, configs, and inventory into NetBox through the Diode pipeline.

BGP sessions are pushed to the netbox-bgp plugin API (not via Diode).

Usage:
    python collectors/network_collector.py --inventory inventory.yaml
    python collectors/network_collector.py --inventory inventory.yaml --dry-run
"""

import argparse
import json
import logging
import os
import re
import sys
from typing import Any

import yaml
from napalm import get_network_driver

from netboxlabs.diode.sdk import DiodeClient, DiodeDryRunClient
from netboxlabs.diode.sdk.ingester import (
    ASN,
    Cable,
    Device,
    DeviceConfig,
    DeviceRole,
    DeviceType,
    Entity,
    GenericObject,
    Interface,
    InventoryItem,
    InventoryItemRole,
    IPAddress,
    Manufacturer,
    Platform,
    Prefix,
    Site,
    VRF,
    VLAN,
    VLANGroup,
)

# Optional: pyATS/Genie for Cisco-specific parsing (CDP, OSPF, IS-IS)
try:
    from genie.conf.base import Device as GenieDevice
    from genie.libs.parser.utils import common as genie_common

    HAS_PYATS = True
except ImportError:
    HAS_PYATS = False

# Optional: requests for pushing data to NetBox plugin APIs
try:
    import requests

    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

log = logging.getLogger("network-collector")

# ---------------------------------------------------------------------------
# Interface type mapping
# ---------------------------------------------------------------------------

SPEED_TO_TYPE = {
    10: "10base-t",
    100: "100base-tx",
    1000: "1000base-t",
    2500: "2.5gbase-t",
    5000: "5gbase-t",
    10000: "10gbase-x-sfpp",
    25000: "25gbase-x-sfp28",
    40000: "40gbase-x-qsfpp",
    50000: "50gbase-x-sfp56",
    100000: "100gbase-x-qsfp28",
}

NAME_TO_TYPE = {
    r"^(Gi|GigabitEthernet)": "1000base-t",
    r"^(Te|TenGigabitEthernet|TenGigE)": "10gbase-x-sfpp",
    r"^(Tw|TwentyFiveGig)": "25gbase-x-sfp28",
    r"^(Fo|FortyGig|FortyGigE)": "40gbase-x-qsfpp",
    r"^(Hu|HundredGig|HundredGigE)": "100gbase-x-qsfp28",
    r"^(Fa|FastEthernet)": "100base-tx",
    r"^(Et|Ethernet)\d": "1000base-t",
    r"^(Lo|Loopback)": "virtual",
    r"^(Vl|Vlan)": "virtual",
    r"^(Tu|Tunnel)": "virtual",
    r"^(Mg|MgmtEth|Management)": "1000base-t",
    r"^(Nu|Null)": "virtual",
    r"^(Po|Port-channel|port-channel)": "lag",
    r"^(BV|BVI)": "bridge",
    r"^(Se|Serial)": "other",
}

DRIVER_TO_PLATFORM = {
    "ios": "Cisco IOS",
    "iosxr": "Cisco IOS-XR",
    "eos": "Arista EOS",
    "junos": "Juniper Junos",
    "nxos": "Cisco NX-OS",
    "nxos_ssh": "Cisco NX-OS",
}

DRIVER_TO_MANUFACTURER = {
    "ios": "Cisco",
    "iosxr": "Cisco",
    "eos": "Arista",
    "junos": "Juniper",
    "nxos": "Cisco",
    "nxos_ssh": "Cisco",
}


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


def load_dotenv(path: str = ".env") -> None:
    """Minimal .env loader — no extra dependency."""
    if not os.path.isfile(path):
        return
    with open(path) as fh:
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, _, val = line.partition("=")
            key = key.strip()
            val = val.strip().strip("\"'")
            os.environ.setdefault(key, val)


def load_inventory(path: str) -> dict:
    """Load the device inventory from a YAML file.

    Expected format:
        defaults:
          site: main
          role: Network Device
          username: admin
          password: cisco
          secret: cisco          # enable secret
          driver: ios            # NAPALM driver
          timeout: 60

        devices:
          - host: 10.10.20.1
            driver: ios
            role: Router
          - host: 10.10.20.55
            driver: ios
            role: Switch
          - host: 10.10.20.100
            driver: iosxr
            role: Router
    """
    with open(path) as fh:
        data = yaml.safe_load(fh)
    if not data or "devices" not in data:
        raise ValueError(f"Inventory file {path} must have a 'devices' list")
    return data


def merge_device_config(device_entry: dict, defaults: dict) -> dict:
    """Merge a single device entry with defaults (device overrides defaults)."""
    merged = dict(defaults)
    merged.update({k: v for k, v in device_entry.items() if v is not None})
    return merged


# ---------------------------------------------------------------------------
# Interface type helpers
# ---------------------------------------------------------------------------


def map_interface_type(name: str, speed: int = 0) -> str:
    """Map an interface name and/or speed to a NetBox interface type."""
    for pattern, iface_type in NAME_TO_TYPE.items():
        if re.match(pattern, name, re.IGNORECASE):
            return iface_type
    if speed and speed in SPEED_TO_TYPE:
        return SPEED_TO_TYPE[speed]
    return "other"


def normalize_interface_name(name: str) -> str:
    """Normalize short interface names from LLDP/CDP to long form for matching."""
    abbrevs = {
        "Gi": "GigabitEthernet",
        "Te": "TenGigabitEthernet",
        "Fa": "FastEthernet",
        "Et": "Ethernet",
        "Fo": "FortyGigabitEthernet",
        "Hu": "HundredGigE",
        "Lo": "Loopback",
        "Vl": "Vlan",
        "Po": "Port-channel",
        "Mg": "MgmtEth",
        "Tu": "Tunnel",
        "Se": "Serial",
    }
    for short, long in abbrevs.items():
        if name.startswith(short) and not name.startswith(long):
            rest = name[len(short):]
            if rest and (rest[0].isdigit() or rest[0] == "/"):
                return long + rest
    return name


# ---------------------------------------------------------------------------
# NAPALM data collection
# ---------------------------------------------------------------------------


def connect_device(host: str, driver: str, username: str, password: str,
                   secret: str = "", timeout: int = 60,
                   optional_args: dict | None = None) -> Any:
    """Connect to a device via NAPALM and return the open driver."""
    driver_cls = get_network_driver(driver)
    opts = optional_args or {}
    if secret:
        opts.setdefault("secret", secret)
    opts.setdefault("transport", "ssh")
    dev = driver_cls(host, username, password, timeout=timeout, optional_args=opts)
    dev.open()
    return dev


def collect_napalm_data(dev: Any) -> dict:
    """Collect all available data from a NAPALM device."""
    data = {}

    # Facts (always available)
    try:
        data["facts"] = dev.get_facts()
        log.info("  Facts: %s (%s)", data["facts"].get("hostname"), data["facts"].get("model"))
    except Exception as exc:
        log.error("  get_facts() failed: %s", exc)
        return data

    # Interfaces
    try:
        data["interfaces"] = dev.get_interfaces()
        log.debug("  Interfaces: %d", len(data["interfaces"]))
    except Exception as exc:
        log.warning("  get_interfaces() failed: %s", exc)

    # Interface IPs (IPv4 + IPv6)
    try:
        data["interfaces_ip"] = dev.get_interfaces_ip()
        log.debug("  Interface IPs: %d interfaces with IPs", len(data["interfaces_ip"]))
    except Exception as exc:
        log.warning("  get_interfaces_ip() failed: %s", exc)

    # LLDP neighbors
    try:
        data["lldp_neighbors"] = dev.get_lldp_neighbors_detail()
        total = sum(len(v) for v in data["lldp_neighbors"].values())
        log.debug("  LLDP neighbors: %d", total)
    except Exception as exc:
        log.debug("  get_lldp_neighbors_detail() failed: %s", exc)

    # Config
    try:
        data["config"] = dev.get_config()
        log.debug("  Config: running=%d bytes", len(data["config"].get("running", "")))
    except Exception as exc:
        log.warning("  get_config() failed: %s", exc)

    # VLANs
    try:
        data["vlans"] = dev.get_vlans()
        log.debug("  VLANs: %d", len(data["vlans"]))
    except Exception as exc:
        log.debug("  get_vlans() unavailable: %s", exc)

    # Network instances (VRFs)
    try:
        data["network_instances"] = dev.get_network_instances()
        log.debug("  VRFs: %d", len(data["network_instances"]))
    except Exception as exc:
        log.debug("  get_network_instances() unavailable: %s", exc)

    # BGP neighbors
    try:
        data["bgp_neighbors"] = dev.get_bgp_neighbors_detail()
        total = sum(len(peers) for vrf in data["bgp_neighbors"].values() for peers in vrf.values())
        log.debug("  BGP peers: %d", total)
    except Exception as exc:
        log.debug("  get_bgp_neighbors_detail() unavailable: %s", exc)

    return data


# ---------------------------------------------------------------------------
# pyATS/Genie data collection (optional)
# ---------------------------------------------------------------------------


def collect_pyats_data(host: str, driver: str, username: str, password: str,
                       secret: str = "") -> dict:
    """Collect CDP, OSPF, IS-IS data via pyATS/Genie parsers."""
    if not HAS_PYATS:
        return {}

    os_map = {"ios": "iosxe", "iosxr": "iosxr", "nxos": "nxos", "nxos_ssh": "nxos"}
    genie_os = os_map.get(driver)
    if not genie_os:
        log.debug("  pyATS: no mapping for driver %s, skipping", driver)
        return {}

    data = {}
    try:
        dev = GenieDevice(name=host, os=genie_os, credentials={
            "default": {"username": username, "password": password}
        })
        dev.connect(ip=host, init_exec_commands=[], init_config_commands=[], log_stdout=False)
    except Exception as exc:
        log.warning("  pyATS connect failed: %s", exc)
        return data

    # CDP neighbors
    try:
        data["cdp_neighbors"] = dev.parse("show cdp neighbors detail")
        log.debug("  CDP neighbors: %d", len(data["cdp_neighbors"].get("index", {})))
    except Exception as exc:
        log.debug("  CDP parse failed: %s", exc)

    # OSPF neighbors
    try:
        data["ospf_neighbors"] = dev.parse("show ip ospf neighbor")
        log.debug("  OSPF neighbors parsed")
    except Exception as exc:
        log.debug("  OSPF parse unavailable: %s", exc)

    # IS-IS adjacencies
    try:
        data["isis_adjacencies"] = dev.parse("show isis adjacency")
        log.debug("  IS-IS adjacencies parsed")
    except Exception as exc:
        log.debug("  IS-IS parse unavailable: %s", exc)

    # Inventory (modules, transceivers)
    try:
        data["inventory"] = dev.parse("show inventory")
        log.debug("  Inventory parsed")
    except Exception as exc:
        log.debug("  Inventory parse unavailable: %s", exc)

    try:
        dev.disconnect()
    except Exception:
        pass

    return data


# ---------------------------------------------------------------------------
# Device reference helper (rich references for Diode reconciler)
# ---------------------------------------------------------------------------


def _device_ref(hostname: str, model: str, manufacturer: str, role: str,
                site_name: str) -> Device:
    """Build a rich Device reference with enough context for the reconciler."""
    return Device(
        name=hostname,
        device_type=DeviceType(
            model=model or "Unknown",
            manufacturer=Manufacturer(name=manufacturer or "Unknown"),
        ),
        role=DeviceRole(name=role),
        site=Site(name=site_name),
    )


# ---------------------------------------------------------------------------
# Entity builders
# ---------------------------------------------------------------------------


def build_device_entity(facts: dict, driver: str, role: str, site_name: str,
                        host: str) -> Entity:
    """Build a Device entity from NAPALM facts."""
    hostname = facts.get("hostname") or host
    model = facts.get("model") or "Unknown"
    vendor = facts.get("vendor") or DRIVER_TO_MANUFACTURER.get(driver, "Unknown")
    serial = facts.get("serial_number") or ""
    os_version = facts.get("os_version") or ""
    platform_name = DRIVER_TO_PLATFORM.get(driver, driver)

    return Entity(device=Device(
        name=hostname,
        device_type=DeviceType(
            model=model,
            manufacturer=Manufacturer(name=vendor),
        ),
        role=DeviceRole(name=role),
        platform=Platform(name=platform_name),
        site=Site(name=site_name),
        serial=serial,
        status="active",
        comments=f"OS: {os_version}" if os_version else "",
        tags=["network-collector"],
    ))


def build_interface_entities(interfaces: dict, hostname: str, model: str,
                             manufacturer: str, role: str,
                             site_name: str) -> list[Entity]:
    """Build Interface entities from NAPALM get_interfaces()."""
    entities = []
    dev_ref = _device_ref(hostname, model, manufacturer, role, site_name)

    for name, iface_data in interfaces.items():
        speed = iface_data.get("speed", 0)
        iface_type = map_interface_type(name, speed)

        iface = Interface(
            device=dev_ref,
            name=name,
            type=iface_type,
            enabled=iface_data.get("is_enabled", True),
            mac_address=iface_data.get("mac_address") or "",
            mtu=iface_data.get("mtu") or 0,
            speed=speed * 1000 if speed else 0,  # NAPALM Mbps → NetBox Kbps
            description=iface_data.get("description") or "",
            tags=["network-collector"],
        )
        entities.append(Entity(interface=iface))

    return entities


def build_ip_entities(interfaces_ip: dict, hostname: str, model: str,
                      manufacturer: str, role: str,
                      site_name: str) -> list[Entity]:
    """Build IPAddress entities from NAPALM get_interfaces_ip()."""
    entities = []
    dev_ref = _device_ref(hostname, model, manufacturer, role, site_name)

    for iface_name, af_data in interfaces_ip.items():
        iface_type = map_interface_type(iface_name)

        for af in ("ipv4", "ipv6"):
            addrs = af_data.get(af, {})
            for addr, meta in addrs.items():
                prefix_len = meta.get("prefix_length", 32 if af == "ipv4" else 128)
                ip_str = f"{addr}/{prefix_len}"

                # Skip link-local IPv6
                if af == "ipv6" and addr.lower().startswith("fe80"):
                    continue

                entities.append(Entity(ip_address=IPAddress(
                    address=ip_str,
                    status="active",
                    assigned_object_interface=Interface(
                        device=dev_ref,
                        name=iface_name,
                        type=iface_type,
                    ),
                    tags=["network-collector"],
                )))

    return entities


def build_vlan_entities(vlans: dict, site_name: str) -> list[Entity]:
    """Build VLAN entities from NAPALM get_vlans()."""
    entities = []
    for vid_str, vlan_data in vlans.items():
        vid = int(vid_str)
        if vid in (0, 4095):
            continue
        name = vlan_data.get("name") or f"VLAN{vid}"
        entities.append(Entity(vlan=VLAN(
            vid=vid,
            name=name,
            site=Site(name=site_name),
            status="active",
            tags=["network-collector"],
        )))
    return entities


def build_vrf_entities(network_instances: dict) -> list[Entity]:
    """Build VRF entities from NAPALM get_network_instances()."""
    entities = []
    for vrf_name, vrf_data in network_instances.items():
        if vrf_name in ("default", "global"):
            continue
        rd = ""
        if vrf_data.get("state", {}).get("route_distinguisher"):
            rd = vrf_data["state"]["route_distinguisher"]
        entities.append(Entity(vrf=VRF(
            name=vrf_name,
            rd=rd or None,
            tags=["network-collector"],
        )))
    return entities


def build_prefix_entities(interfaces_ip: dict, site_name: str) -> list[Entity]:
    """Build Prefix entities from discovered interface IPs."""
    import ipaddress
    seen = set()
    entities = []

    for iface_name, af_data in interfaces_ip.items():
        for af in ("ipv4", "ipv6"):
            for addr, meta in af_data.get(af, {}).items():
                prefix_len = meta.get("prefix_length", 32 if af == "ipv4" else 128)
                try:
                    network = ipaddress.ip_network(f"{addr}/{prefix_len}", strict=False)
                except ValueError:
                    continue
                # Skip host routes and link-local
                if af == "ipv4" and prefix_len >= 31:
                    continue
                if af == "ipv6" and (prefix_len >= 127 or addr.lower().startswith("fe80")):
                    continue
                prefix_str = str(network)
                if prefix_str not in seen:
                    seen.add(prefix_str)
                    entities.append(Entity(prefix=Prefix(
                        prefix=prefix_str,
                        scope_site=Site(name=site_name),
                        status="active",
                        tags=["network-collector"],
                    )))

    return entities


def build_config_entity(config: dict, hostname: str, model: str,
                        manufacturer: str, role: str,
                        site_name: str) -> Entity | None:
    """Build a DeviceConfig entity from NAPALM get_config()."""
    running = config.get("running", "")
    startup = config.get("startup", "")
    if not running and not startup:
        return None
    dev_ref = _device_ref(hostname, model, manufacturer, role, site_name)
    return Entity(device=dev_ref, device_config=DeviceConfig(
        running=running.encode("utf-8") if running else None,
        startup=startup.encode("utf-8") if startup else None,
    ))


def build_inventory_entities(inventory: dict, hostname: str, model: str,
                             manufacturer: str, role: str,
                             site_name: str) -> list[Entity]:
    """Build InventoryItem entities from pyATS show inventory."""
    entities = []
    dev_ref = _device_ref(hostname, model, manufacturer, role, site_name)

    # pyATS inventory format varies by platform; handle the common structure
    items = inventory.get("main", {}).get("chassis", {})
    if not items:
        items = inventory
    for item_name, item_data in items.items():
        if isinstance(item_data, dict):
            pid = item_data.get("pid") or item_data.get("name") or ""
            sn = item_data.get("sn") or ""
            desc = item_data.get("descr") or item_data.get("description") or ""
            entities.append(Entity(inventory_item=InventoryItem(
                device=dev_ref,
                name=str(item_name)[:64],
                part_id=str(pid)[:50] if pid else None,
                serial=str(sn)[:50] if sn else None,
                description=str(desc)[:200] if desc else None,
                discovered=True,
                tags=["network-collector"],
            )))

    return entities


# ---------------------------------------------------------------------------
# Cable discovery from LLDP/CDP
# ---------------------------------------------------------------------------


def build_cable_entities_from_lldp(
    lldp_all: dict[str, dict],
    site_name: str,
    device_models: dict[str, tuple[str, str, str]],
) -> list[Entity]:
    """Build Cable entities from LLDP neighbor data collected from all devices.

    lldp_all: {hostname: {local_iface: [{remote_system_name, remote_port, ...}]}}
    device_models: {hostname: (model, manufacturer, role)}

    Deduplication: each link appears on both ends. We sort the pair and keep
    only one Cable per unique (deviceA:portA, deviceB:portB) tuple.
    """
    seen_links: set[tuple] = set()
    entities = []

    for local_host, iface_neighbors in lldp_all.items():
        local_info = device_models.get(local_host, ("Unknown", "Unknown", "Network Device"))

        for local_iface, neighbors in iface_neighbors.items():
            for neighbor in neighbors:
                remote_host = (
                    neighbor.get("remote_system_name")
                    or neighbor.get("remote_system_description", "")
                ).split(".")[0]  # Strip FQDN
                remote_port = neighbor.get("remote_port") or neighbor.get("remote_port_description", "")

                if not remote_host or not remote_port:
                    continue

                # Normalize interface names for matching
                local_norm = normalize_interface_name(local_iface)
                remote_norm = normalize_interface_name(remote_port)

                # Deduplicate: sorted link key
                link_a = f"{local_host}:{local_norm}"
                link_b = f"{remote_host}:{remote_norm}"
                link_key = tuple(sorted([link_a, link_b]))
                if link_key in seen_links:
                    continue
                seen_links.add(link_key)

                # Build device references for both ends
                remote_info = device_models.get(
                    remote_host, ("Unknown", "Unknown", "Network Device")
                )

                a_dev = _device_ref(local_host, *local_info, site_name)
                b_dev = _device_ref(remote_host, *remote_info, site_name)

                a_iface_type = map_interface_type(local_norm)
                b_iface_type = map_interface_type(remote_norm)

                cable = Cable(
                    a_terminations=[GenericObject(object_interface=Interface(
                        device=a_dev,
                        name=local_norm,
                        type=a_iface_type,
                    ))],
                    b_terminations=[GenericObject(object_interface=Interface(
                        device=b_dev,
                        name=remote_norm,
                        type=b_iface_type,
                    ))],
                    status="connected",
                    tags=["lldp-discovered"],
                )
                entities.append(Entity(cable=cable))
                log.info("  Cable: %s:%s <-> %s:%s",
                         local_host, local_norm, remote_host, remote_norm)

    return entities


def build_cable_entities_from_cdp(
    cdp_all: dict[str, dict],
    site_name: str,
    device_models: dict[str, tuple[str, str, str]],
    existing_links: set[tuple] | None = None,
) -> list[Entity]:
    """Build Cable entities from CDP neighbor data (pyATS parsed).

    cdp_all: {hostname: pyats_parsed_cdp_output}
    Only creates cables for links NOT already discovered via LLDP.
    """
    if existing_links is None:
        existing_links = set()
    entities = []

    for local_host, cdp_data in cdp_all.items():
        local_info = device_models.get(local_host, ("Unknown", "Unknown", "Network Device"))

        # pyATS CDP format: {"index": {1: {device_id, local_interface, port_id, ...}}}
        for idx, entry in cdp_data.get("index", {}).items():
            remote_host = (entry.get("device_id") or "").split(".")[0]
            local_iface = entry.get("local_interface") or ""
            remote_port = entry.get("port_id") or ""

            if not remote_host or not local_iface or not remote_port:
                continue

            local_norm = normalize_interface_name(local_iface)
            remote_norm = normalize_interface_name(remote_port)

            link_a = f"{local_host}:{local_norm}"
            link_b = f"{remote_host}:{remote_norm}"
            link_key = tuple(sorted([link_a, link_b]))
            if link_key in existing_links:
                continue
            existing_links.add(link_key)

            remote_info = device_models.get(
                remote_host, ("Unknown", "Unknown", "Network Device")
            )

            a_dev = _device_ref(local_host, *local_info, site_name)
            b_dev = _device_ref(remote_host, *remote_info, site_name)

            cable = Cable(
                a_terminations=[GenericObject(object_interface=Interface(
                    device=a_dev,
                    name=local_norm,
                    type=map_interface_type(local_norm),
                ))],
                b_terminations=[GenericObject(object_interface=Interface(
                    device=b_dev,
                    name=remote_norm,
                    type=map_interface_type(remote_norm),
                ))],
                status="connected",
                tags=["cdp-discovered"],
            )
            entities.append(Entity(cable=cable))
            log.info("  Cable (CDP): %s:%s <-> %s:%s",
                     local_host, local_norm, remote_host, remote_norm)

    return entities


# ---------------------------------------------------------------------------
# NetBox plugin API pushers (BGP, OSPF, IS-IS)
# ---------------------------------------------------------------------------


def push_bgp_to_netbox(bgp_data: dict, hostname: str, netbox_url: str,
                       netbox_token: str, site_name: str) -> None:
    """Push BGP neighbor data to the netbox-bgp plugin API."""
    if not HAS_REQUESTS or not bgp_data:
        return

    headers = {"Authorization": f"Bearer {netbox_token}", "Content-Type": "application/json"}
    base = netbox_url.rstrip("/")

    # First look up the local device in NetBox to get its ID
    resp = requests.get(f"{base}/api/dcim/devices/", params={"name": hostname},
                        headers=headers, timeout=10)
    if resp.status_code != 200:
        log.warning("  BGP push: cannot find device %s in NetBox", hostname)
        return
    devices = resp.json().get("results", [])
    if not devices:
        log.warning("  BGP push: device %s not found in NetBox", hostname)
        return
    local_device_id = devices[0]["id"]

    for vrf_name, peers in bgp_data.items():
        for peer_ip, peer_list in peers.items():
            for peer in peer_list:
                local_as = peer.get("local_as")
                remote_as = peer.get("remote_as")
                remote_id = peer.get("remote_id", "")

                if not local_as or not remote_as:
                    continue

                session_data = {
                    "name": f"{hostname} <-> {peer_ip}",
                    "device": local_device_id,
                    "local_address": None,  # Would need to resolve
                    "remote_address": None,
                    "local_as": {"asn": local_as},
                    "remote_as": {"asn": remote_as},
                    "status": "active" if peer.get("is_enabled") else "offline",
                    "description": f"Auto-discovered by network-collector",
                }

                try:
                    resp = requests.post(
                        f"{base}/api/plugins/bgp/sessions/",
                        headers=headers,
                        json=session_data,
                        timeout=10,
                    )
                    if resp.status_code in (200, 201):
                        log.info("  BGP session created: %s AS%s <-> %s AS%s",
                                 hostname, local_as, peer_ip, remote_as)
                    elif resp.status_code == 400 and "already exists" in resp.text.lower():
                        log.debug("  BGP session already exists: %s <-> %s", hostname, peer_ip)
                    else:
                        log.warning("  BGP push failed (%d): %s", resp.status_code, resp.text[:200])
                except Exception as exc:
                    log.warning("  BGP push error: %s", exc)


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------


def collect_all_entities(inventory: dict, env_file: str = ".env") -> tuple[list[Entity], dict]:
    """Walk all devices in inventory, collect data, build entities."""
    defaults = inventory.get("defaults", {})
    site_name = defaults.get("site", "main")
    entities: list[Entity] = []
    lldp_all: dict[str, dict] = {}
    cdp_all: dict[str, dict] = {}
    device_models: dict[str, tuple[str, str, str]] = {}
    bgp_all: dict[str, dict] = {}

    for dev_entry in inventory["devices"]:
        cfg = merge_device_config(dev_entry, defaults)
        host = cfg["host"]
        driver = cfg.get("driver", "ios")
        role = cfg.get("role", "Network Device")
        username = cfg.get("username", "admin")
        password = cfg.get("password", "")
        secret = cfg.get("secret", "")
        timeout = int(cfg.get("timeout", 60))
        optional_args = cfg.get("optional_args", {})

        log.info("Connecting to %s (driver=%s, role=%s)...", host, driver, role)

        # --- NAPALM collection ---
        try:
            dev = connect_device(host, driver, username, password, secret,
                                 timeout, optional_args)
        except Exception as exc:
            log.error("Failed to connect to %s: %s", host, exc)
            continue

        try:
            napalm_data = collect_napalm_data(dev)
        finally:
            try:
                dev.close()
            except Exception:
                pass

        if not napalm_data.get("facts"):
            log.error("No facts for %s, skipping", host)
            continue

        facts = napalm_data["facts"]
        hostname = facts.get("hostname") or host
        model = facts.get("model") or "Unknown"
        vendor = facts.get("vendor") or DRIVER_TO_MANUFACTURER.get(driver, "Unknown")

        # Track device info for cable building
        device_models[hostname] = (model, vendor, role)

        # Device entity
        entities.append(build_device_entity(facts, driver, role, site_name, host))

        # Interface entities
        if napalm_data.get("interfaces"):
            entities.extend(build_interface_entities(
                napalm_data["interfaces"], hostname, model, vendor, role, site_name
            ))

        # IP entities (IPv4 + IPv6)
        if napalm_data.get("interfaces_ip"):
            entities.extend(build_ip_entities(
                napalm_data["interfaces_ip"], hostname, model, vendor, role, site_name
            ))
            # Prefix entities from discovered IPs
            entities.extend(build_prefix_entities(
                napalm_data["interfaces_ip"], site_name
            ))

        # VLAN entities
        if napalm_data.get("vlans"):
            entities.extend(build_vlan_entities(napalm_data["vlans"], site_name))

        # VRF entities
        if napalm_data.get("network_instances"):
            entities.extend(build_vrf_entities(napalm_data["network_instances"]))

        # Config entity
        if napalm_data.get("config"):
            config_entity = build_config_entity(
                napalm_data["config"], hostname, model, vendor, role, site_name
            )
            if config_entity:
                entities.append(config_entity)

        # LLDP neighbors (saved for cable building later)
        if napalm_data.get("lldp_neighbors"):
            lldp_all[hostname] = napalm_data["lldp_neighbors"]

        # BGP data (saved for plugin API push later)
        if napalm_data.get("bgp_neighbors"):
            bgp_all[hostname] = napalm_data["bgp_neighbors"]

        # --- pyATS collection (optional) ---
        if HAS_PYATS and driver in ("ios", "iosxr", "nxos", "nxos_ssh"):
            log.info("  Running pyATS parsers...")
            pyats_data = collect_pyats_data(host, driver, username, password, secret)

            if pyats_data.get("cdp_neighbors"):
                cdp_all[hostname] = pyats_data["cdp_neighbors"]

            if pyats_data.get("inventory"):
                entities.extend(build_inventory_entities(
                    pyats_data["inventory"], hostname, model, vendor, role, site_name
                ))

    # --- Cable entities from LLDP ---
    if lldp_all:
        log.info("Building cable entities from LLDP data...")
        cable_entities = build_cable_entities_from_lldp(lldp_all, site_name, device_models)
        entities.extend(cable_entities)

        # Extract seen links for CDP dedup
        seen_links = set()
        for local_host, iface_neighbors in lldp_all.items():
            for local_iface, neighbors in iface_neighbors.items():
                for neighbor in neighbors:
                    remote_host = (neighbor.get("remote_system_name") or "").split(".")[0]
                    remote_port = neighbor.get("remote_port") or ""
                    if remote_host and remote_port:
                        local_norm = normalize_interface_name(local_iface)
                        remote_norm = normalize_interface_name(remote_port)
                        link_key = tuple(sorted([
                            f"{local_host}:{local_norm}",
                            f"{remote_host}:{remote_norm}",
                        ]))
                        seen_links.add(link_key)
    else:
        seen_links = set()

    # --- Cable entities from CDP (only new links) ---
    if cdp_all:
        log.info("Building cable entities from CDP data...")
        cdp_cable_entities = build_cable_entities_from_cdp(
            cdp_all, site_name, device_models, seen_links
        )
        entities.extend(cdp_cable_entities)

    return entities, bgp_all


def ingest_entities(entities: list[Entity], dry_run: bool = False) -> None:
    """Send entities to Diode (or dry-run print them)."""
    if not entities:
        log.warning("No entities to ingest")
        return

    target = os.environ.get("DIODE_TARGET", "grpc://localhost:8080/diode")
    client_id = os.environ.get("DIODE_CLIENT_ID",
                               os.environ.get("INGESTER_CLIENT_ID", "diode-ingester"))
    client_secret = os.environ.get("DIODE_CLIENT_SECRET",
                                   os.environ.get("INGESTER_CLIENT_SECRET", ""))

    if dry_run:
        log.info("DRY RUN: %d entities would be ingested", len(entities))
        for i, e in enumerate(entities):
            log.info("  [%d] %s", i, e)
        return

    if not client_secret:
        log.error("DIODE_CLIENT_SECRET not set — cannot ingest")
        sys.exit(1)

    log.info("Ingesting %d entities to %s ...", len(entities), target)

    from netboxlabs.diode.sdk.ingester import create_message_chunks

    with DiodeClient(
        target=target,
        client_id=client_id,
        client_secret=client_secret,
        app_name="network-collector",
        app_version="0.1.0",
    ) as client:
        chunks = create_message_chunks(entities)
        for idx, chunk in enumerate(chunks):
            resp = client.ingest(entities=chunk)
            if resp.errors:
                log.error("Chunk %d errors: %s", idx, resp.errors)
            else:
                log.info("Chunk %d: %d entities ingested", idx, len(chunk))


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(description="Network device collector for NetBox")
    parser.add_argument("--inventory", "-i", required=True,
                        help="Path to device inventory YAML file")
    parser.add_argument("--dry-run", action="store_true",
                        help="Collect data but don't ingest")
    parser.add_argument("--log-level", default="INFO",
                        choices=["DEBUG", "INFO", "WARNING", "ERROR"])
    parser.add_argument("--env-file", default=".env",
                        help="Path to .env file (default: .env)")
    parser.add_argument("--no-bgp-push", action="store_true",
                        help="Skip pushing BGP data to netbox-bgp plugin")
    parser.add_argument("--no-pyats", action="store_true",
                        help="Skip pyATS/Genie collection even if available")
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s %(name)s %(levelname)s %(message)s",
    )

    load_dotenv(args.env_file)

    if args.no_pyats:
        global HAS_PYATS
        HAS_PYATS = False

    # Load inventory
    inventory = load_inventory(args.inventory)
    log.info("Loaded %d devices from inventory", len(inventory["devices"]))

    # Collect and build entities
    entities, bgp_all = collect_all_entities(inventory, args.env_file)
    log.info("Total entities: %d", len(entities))

    # Ingest via Diode
    ingest_entities(entities, dry_run=args.dry_run)

    # Push BGP sessions to netbox-bgp plugin API
    if not args.no_bgp_push and not args.dry_run and bgp_all:
        netbox_url = os.environ.get("NETBOX_URL", "http://172.19.77.160:8000")
        netbox_token = os.environ.get("NETBOX_API_TOKEN", "")
        if netbox_token:
            for hostname, bgp_data in bgp_all.items():
                log.info("Pushing BGP data for %s to netbox-bgp...", hostname)
                push_bgp_to_netbox(bgp_data, hostname, netbox_url, netbox_token,
                                   inventory.get("defaults", {}).get("site", "main"))
        else:
            log.warning("NETBOX_API_TOKEN not set, skipping BGP plugin push")

    log.info("Done!")


if __name__ == "__main__":
    main()