lldp-mapper/exports.py

# exports.py - Generate CSV, Mermaid, and Graphviz outputs from DB
import csv
import os
import logging
from config import EXPORTS_DIR
from db import get_all_switches, get_all_links

logger = logging.getLogger(__name__)


def _ensure_exports_dir():
    os.makedirs(EXPORTS_DIR, exist_ok=True)


def export_csv():
    """Export links + switch info to CSV."""
    _ensure_exports_dir()
    path = os.path.join(EXPORTS_DIR, "topology.csv")
    switches = {s['chassis_id']: s for s in get_all_switches()}
    links = get_all_links()

    with open(path, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow([
            'Switch A Hostname', 'Switch A IP', 'Switch A Chassis',
            'Port A',
            'Switch B Hostname', 'Switch B IP', 'Switch B Chassis',
            'Port B'
        ])
        for link in links:
            sw_a = switches.get(link['chassis_a'], {})
            sw_b = switches.get(link['chassis_b'], {})
            writer.writerow([
                sw_a.get('hostname', link['chassis_a']),
                sw_a.get('mgmt_ip', ''),
                link['chassis_a'],
                link['port_a'],
                sw_b.get('hostname', link['chassis_b']),
                sw_b.get('mgmt_ip', ''),
                link['chassis_b'],
                link['port_b'],
            ])

    logger.info(f"CSV exported to {path}")
    return path


def export_mermaid():
    """Export topology as clean Mermaid diagram with reference tables."""
    _ensure_exports_dir()
    path = os.path.join(EXPORTS_DIR, "topology.md")
    switches = {s['chassis_id']: s for s in get_all_switches()}
    links = get_all_links()

    def node_id(sw, chassis_id):
        return sw.get('hostname', chassis_id).replace('-', '_').replace('.', '_')

    lines = [
        "# Network Topology",
        "",
        "```mermaid",
        "graph LR",
    ]

    # Node definitions - hostname only as label
    for chassis_id, sw in sorted(switches.items(), key=lambda x: x[1].get('hostname','')):
        nid = node_id(sw, chassis_id)
        hostname = sw.get('hostname', chassis_id)
        lines.append(f'    {nid}["{hostname}"]')

    lines.append("")

    # Edge definitions - portA --> portB as label
    for link in links:
        sw_a = switches.get(link['chassis_a'], {})
        sw_b = switches.get(link['chassis_b'], {})
        id_a = node_id(sw_a, link['chassis_a'])
        id_b = node_id(sw_b, link['chassis_b'])
        lines.append(f'    {id_a} -- "{link["port_a"]} to {link["port_b"]}" --> {id_b}')

    lines.append("```")
    lines.append("")

    # Switch reference table
    lines.append("## Switch Reference")
    lines.append("")
    lines.append("| Hostname | Management IP | Chassis ID | Last Seen |")
    lines.append("|----------|---------------|------------|-----------|")
    for sw in sorted(switches.values(), key=lambda x: x.get('hostname', '')):
        lines.append(
            f"| {sw.get('hostname','')} "
            f"| {sw.get('mgmt_ip','')} "
            f"| {sw.get('chassis_id','')} "
            f"| {sw.get('last_seen','')} |"
        )

    lines.append("")

    # Link reference table
    lines.append("## Link Reference")
    lines.append("")
    lines.append("| Switch A | Port A | Switch B | Port B |")
    lines.append("|----------|--------|----------|--------|")
    for link in links:
        sw_a = switches.get(link['chassis_a'], {})
        sw_b = switches.get(link['chassis_b'], {})
        lines.append(
            f"| {sw_a.get('hostname', link['chassis_a'])} "
            f"| {link['port_a']} "
            f"| {sw_b.get('hostname', link['chassis_b'])} "
            f"| {link['port_b']} |"
        )

    with open(path, 'w') as f:
        f.write('\n'.join(lines))

    logger.info(f"Mermaid exported to {path}")
    return path


def export_graphviz():
    """Export topology as Graphviz DOT and render to PNG."""
    _ensure_exports_dir()
    dot_path = os.path.join(EXPORTS_DIR, "topology.dot")
    png_path = os.path.join(EXPORTS_DIR, "topology.png")

    switches = {s['chassis_id']: s for s in get_all_switches()}
    links = get_all_links()

    def nid(chassis):
        sw = switches.get(chassis, {})
        return sw.get('hostname', chassis).replace('-', '_').replace('.', '_')

    dot = []
    dot.append('digraph network {')
    dot.append('    rankdir=LR;')
    dot.append('    bgcolor="white";')
    dot.append('    pad=0.8;')
    dot.append('    nodesep=0.6;')
    dot.append('    ranksep=1.5;')
    dot.append('    node [shape=box, style="filled,rounded", fillcolor="#1a3a6e", fontcolor=white, fontname="Helvetica Bold", fontsize=12, margin="0.3,0.2", width=2.2];')
    dot.append('    edge [fontname="Helvetica", fontsize=9, fontcolor="#444444", color="#1a3a6e", penwidth=1.5, dir=none];')
    dot.append('')

    for chassis_id, sw in switches.items():
        n = nid(chassis_id)
        hostname = sw.get('hostname', chassis_id)
        mgmt_ip = sw.get('mgmt_ip', '')
        label = f"{hostname}\n{mgmt_ip}" if mgmt_ip else hostname
        dot.append(f'    {n} [label="{label}"];')

    dot.append('')

    for link in links:
        a = nid(link["chassis_a"])
        b = nid(link["chassis_b"])
        pa = link["port_a"]
        pb = link["port_b"]
        dot.append(f'    {a} -> {b} [taillabel="{pa}", headlabel="{pb}"];')

    dot.append('}')

    with open(dot_path, 'w') as f:
        f.write('\n'.join(dot))

    try:
        import subprocess
        result = subprocess.run(
            ['dot', '-Tpng', '-Gdpi=150', dot_path, '-o', png_path],
            capture_output=True, text=True, timeout=30
        )
        if result.returncode == 0:
            logger.info(f"Graphviz PNG exported to {png_path}")
            return png_path
        else:
            logger.error(f"Graphviz render failed: {result.stderr}")
            return dot_path
    except FileNotFoundError:
        logger.warning("graphviz not found")
        return dot_path
    except Exception as e:
        logger.error(f"Graphviz error: {e}")
        return dot_path


def run_all_exports():
    csv_path = export_csv()
    md_path = export_mermaid()
    png_path = export_graphviz()
    return {"csv": csv_path, "mermaid": md_path, "graphviz": png_path}