lldp-mapper/db.py

# db.py - SQLite database operations
import sqlite3
import re
import os
import logging
from config import DB_PATH, SSH_USERNAME, SSH_PASSWORD

logger = logging.getLogger(__name__)


def get_conn():
    conn = sqlite3.connect(DB_PATH, check_same_thread=False, timeout=10)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = get_conn()
    c = conn.cursor()

    c.execute("""
        CREATE TABLE IF NOT EXISTS switches (
            chassis_id   TEXT PRIMARY KEY,
            hostname     TEXT,
            mgmt_ip      TEXT,
            description  TEXT,
            firmware     TEXT DEFAULT '',
            vendor       TEXT DEFAULT '',
            last_seen    TEXT
        )
    """)
    # Migrate existing DBs that predate firmware/vendor/serial columns
    for col, default in [('firmware', ''), ('vendor', ''), ('serial', '')]:
        try:
            c.execute(f"ALTER TABLE switches ADD COLUMN {col} TEXT DEFAULT '{default}'")
        except Exception:
            pass  # Column already exists

    c.execute("""
        CREATE TABLE IF NOT EXISTS links (
            id              INTEGER PRIMARY KEY AUTOINCREMENT,
            chassis_a       TEXT,
            port_a          TEXT,
            chassis_b       TEXT,
            port_b          TEXT,
            UNIQUE(chassis_a, port_a, chassis_b, port_b)
        )
    """)

    c.execute("""
        CREATE TABLE IF NOT EXISTS settings (
            key   TEXT PRIMARY KEY,
            value TEXT
        )
    """)

    c.execute("""
        CREATE TABLE IF NOT EXISTS scan_log (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            started_at  TEXT,
            finished_at TEXT,
            status      TEXT,
            switches_ok INTEGER DEFAULT 0,
            switches_fail INTEGER DEFAULT 0
        )
    """)

    c.execute("INSERT OR IGNORE INTO settings VALUES ('autoscan_enabled', 'false')")
    c.execute("INSERT OR IGNORE INTO settings VALUES ('autoscan_interval', '60')")
    # Seed SSH credentials from config.py on first run only — the settings
    # table is the source of truth from then on (editable via /api/credentials).
    c.execute("INSERT OR IGNORE INTO settings (key, value) VALUES ('ssh_username', ?)", (SSH_USERNAME,))
    c.execute("INSERT OR IGNORE INTO settings (key, value) VALUES ('ssh_password', ?)", (SSH_PASSWORD,))

    c.execute(
        "CREATE TABLE IF NOT EXISTS node_positions "
        "(chassis_id TEXT PRIMARY KEY, x REAL, y REAL)"
    )

    conn.commit()
    conn.close()


def upsert_switch(chassis_id, hostname, mgmt_ip, description, firmware='', vendor='', serial='', trusted=False):
    """
    trusted=True means mgmt_ip came from directly scanning this switch (its own
    SSH session / local LLDP info) — always authoritative. trusted=False means
    it came from another switch's LLDP neighbor report about this one: some
    switches (e.g. HP/Aruba ProCurve) advertise a per-port Management Address
    TLV that reflects whatever VLAN happens to be routable out that specific
    link, not the switch's real management IP — so an untrusted mgmt_ip is
    only used to fill in a still-unknown value, never to overwrite one already
    established by a trusted (direct) scan.
    """
    conn = get_conn()
    conn.execute("""
        INSERT INTO switches (chassis_id, hostname, mgmt_ip, description, firmware, vendor, serial, last_seen)
        VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'))
        ON CONFLICT(chassis_id) DO UPDATE SET
            hostname    = excluded.hostname,
            mgmt_ip     = CASE
                               WHEN excluded.mgmt_ip = '' OR excluded.mgmt_ip NOT LIKE '%.%.%.%' THEN mgmt_ip
                               WHEN ? THEN excluded.mgmt_ip
                               WHEN mgmt_ip = '' OR mgmt_ip IS NULL THEN excluded.mgmt_ip
                               ELSE mgmt_ip
                           END,
            description = CASE WHEN excluded.description != '' THEN excluded.description ELSE description END,
            firmware    = CASE WHEN excluded.firmware    != '' THEN excluded.firmware    ELSE firmware    END,
            vendor      = CASE WHEN excluded.vendor      != '' THEN excluded.vendor      ELSE vendor      END,
            serial      = CASE WHEN excluded.serial      != '' THEN excluded.serial      ELSE serial      END,
            last_seen   = excluded.last_seen
    """, (chassis_id, hostname, mgmt_ip, description, firmware, vendor, serial, bool(trusted)))
    conn.commit()
    conn.close()


def _port_num(port):
    """Trailing port digit(s): 'Gi1/9' → '9', 'Te1/10' → '10', '9' → '9'."""
    m = re.search(r'(\d+)$', (port or '').strip())
    return m.group(1) if m else (port or '').strip()


def upsert_link(chassis_a, port_a, chassis_b, port_b):
    """
    chassis_a/port_a is always the scanned switch's own self-reported port —
    trustworthy. chassis_b/port_b is that switch's perception of its neighbor's
    port, which for LLDP is often the neighbor's raw wire-transmitted Port ID
    (an ifIndex-style number) rather than the neighbor's own friendly port name.
    When the neighbor later gets scanned directly, its own upsert_link call
    supplies the trustworthy label for what is now the "chassis_b" side here —
    so on a match we overwrite that side with the newly-known self-report
    instead of leaving the untrusted guess in place.
    """
    self_chassis, self_port = chassis_a, port_a

    # Normalize order so A→B and B→A are treated as the same link
    if chassis_a > chassis_b:
        chassis_a, chassis_b = chassis_b, chassis_a
        port_a, port_b = port_b, port_a

    conn = get_conn()
    # Exact match first
    existing = conn.execute("""
        SELECT id FROM links WHERE chassis_a=? AND port_a=? AND chassis_b=? AND port_b=?
    """, (chassis_a, port_a, chassis_b, port_b)).fetchone()
    if not existing:
        # Fuzzy match: same chassis pair, same trailing port numbers
        # handles 'Gi1/9' vs '9' reported by different vendors for the same cable
        rows = conn.execute("""
            SELECT id, port_a, port_b FROM links WHERE chassis_a=? AND chassis_b=?
        """, (chassis_a, chassis_b)).fetchall()
        for r in rows:
            if _port_num(port_a) == _port_num(r[1]) and _port_num(port_b) == _port_num(r[2]):
                existing = r
                break
            # A single physical port can only carry one live link, so an exact
            # match on either side's port name (e.g. Aruba's front-panel label
            # 'E5' vs the ifIndex-style '101' its neighbors report for it) means
            # this is the same cable even though the trailing digits don't align.
            if port_a == r[1] or port_b == r[2]:
                existing = r
                break

    if existing:
        # Overwrite whichever stored side corresponds to the self-reporting
        # switch with its trusted value — it may currently hold the neighbor's
        # untrusted guess from when this link was first inserted.
        col = 'port_a' if self_chassis == chassis_a else 'port_b'
        conn.execute(f"UPDATE links SET {col}=? WHERE id=?", (self_port, existing['id']))
        conn.commit()
    else:
        conn.execute("""
            INSERT OR IGNORE INTO links (chassis_a, port_a, chassis_b, port_b)
            VALUES (?, ?, ?, ?)
        """, (chassis_a, port_a, chassis_b, port_b))
        conn.commit()
    conn.close()


def prune_stale_links(chassis_id, current):
    """
    Drop links on a just-scanned switch that its own LLDP table no longer backs.
    current = set of (local_port, neighbor_chassis_id) the switch reported this
    scan. Dept scans don't clear_links() first, so without this a re-cabled port
    (e.g. a new switch inserted mid-run) keeps its old link alongside the new one.
    Must run after the switch's upsert_link calls, which rewrite its side of each
    surviving link to the self-reported port name compared here.
    """
    conn = get_conn()
    rows = conn.execute("""
        SELECT id, chassis_a, port_a, chassis_b, port_b FROM links
        WHERE chassis_a=? OR chassis_b=?
    """, (chassis_id, chassis_id)).fetchall()
    stale = []
    for r in rows:
        if r['chassis_a'] == chassis_id:
            key = (r['port_a'], r['chassis_b'])
        else:
            key = (r['port_b'], r['chassis_a'])
        if key not in current:
            stale.append((r['id'],))
    conn.executemany("DELETE FROM links WHERE id=?", stale)
    conn.commit()
    conn.close()
    if stale:
        logger.info(f"prune_stale_links: removed {len(stale)} stale link(s) on {chassis_id}")
    return len(stale)


def clear_stale_switches(since_dt):
    """Remove switches not seen since the given datetime string."""
    conn = get_conn()
    conn.execute("DELETE FROM switches WHERE last_seen < ? OR last_seen IS NULL", (since_dt,))
    deleted = conn.execute("SELECT changes()").fetchone()[0]
    conn.commit()
    conn.close()
    return deleted

def clear_links():
    conn = get_conn()
    conn.execute("DELETE FROM links")
    conn.commit()
    conn.close()


def _is_mac(chassis):
    return bool(re.match(r'^[0-9A-Fa-f]{2}[-:][0-9A-Fa-f]{2}', chassis))


def _is_real_ip(val):
    return bool(re.match(r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$', val or ''))


def merge_duplicate_switches():
    """
    After a scan, some switches may have two records:
      - One with MAC chassis_id (seen as neighbor by another switch)
      - One with IP chassis_id (directly scanned but chassis extraction failed)
    This function merges them: keeps MAC as canonical, fixes links, removes IP record.
    """
    conn = get_conn()

    dupes = conn.execute("""
        SELECT hostname, COUNT(*) as cnt, GROUP_CONCAT(chassis_id) as ids,
               GROUP_CONCAT(mgmt_ip) as ips
        FROM switches
        GROUP BY hostname
        HAVING cnt > 1
    """).fetchall()

    merged = 0
    for row in dupes:
        ids = row['ids'].split(',')
        ips_raw = row['ips'].split(',')

        mac_ids = [i for i in ids if _is_mac(i)]
        ip_ids  = [i for i in ids if not _is_mac(i)]
        real_ips = [v for v in ips_raw if _is_real_ip(v)]

        if not mac_ids or not ip_ids:
            continue

        canonical = mac_ids[0]
        best_ip   = real_ips[0] if real_ips else ''

        # Update canonical with best mgmt_ip
        if best_ip:
            conn.execute("UPDATE switches SET mgmt_ip=? WHERE chassis_id=?",
                         (best_ip, canonical))

        # Repoint links
        for old_id in ip_ids:
            conn.execute("UPDATE links SET chassis_a=? WHERE chassis_a=?", (canonical, old_id))
            conn.execute("UPDATE links SET chassis_b=? WHERE chassis_b=?", (canonical, old_id))
            conn.execute("DELETE FROM switches WHERE chassis_id=?", (old_id,))

        merged += 1
        logger.info(f"Merged {row['hostname']}: {ip_ids} -> {canonical} (mgmt_ip={best_ip})")

    # Remove any self-links or exact duplicate links created by repointing
    # Remove duplicate links - normalize chassis+port order before grouping
    conn.execute("""
        DELETE FROM links WHERE id NOT IN (
            SELECT MIN(id) FROM links
            GROUP BY
                CASE WHEN chassis_a < chassis_b THEN chassis_a ELSE chassis_b END,
                CASE WHEN chassis_a < chassis_b THEN chassis_b ELSE chassis_a END,
                CASE WHEN chassis_a < chassis_b THEN port_a ELSE port_b END,
                CASE WHEN chassis_a < chassis_b THEN port_b ELSE port_a END
        )
    """)
    conn.execute("DELETE FROM links WHERE chassis_a = chassis_b")

    conn.commit()
    conn.close()
    logger.info(f"merge_duplicate_switches: merged {merged} hostname groups")


def get_all_switches():
    conn = get_conn()
    rows = conn.execute("SELECT * FROM switches ORDER BY hostname").fetchall()
    conn.close()
    return [dict(r) for r in rows]


def get_all_links():
    conn = get_conn()
    rows = conn.execute("""
        SELECT l.*,
               sa.hostname as hostname_a, sa.mgmt_ip as ip_a,
               sb.hostname as hostname_b, sb.mgmt_ip as ip_b
        FROM links l
        LEFT JOIN switches sa ON l.chassis_a = sa.chassis_id
        LEFT JOIN switches sb ON l.chassis_b = sb.chassis_id
    """).fetchall()
    conn.close()
    return [dict(r) for r in rows]


def get_setting(key):
    conn = get_conn()
    row = conn.execute("SELECT value FROM settings WHERE key=?", (key,)).fetchone()
    conn.close()
    return row["value"] if row else None


def set_setting(key, value):
    conn = get_conn()
    conn.execute("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)", (key, str(value)))
    conn.commit()
    conn.close()


def log_scan_start():
    conn = get_conn()
    c = conn.cursor()
    c.execute("INSERT INTO scan_log (started_at, status) VALUES (datetime('now'), 'running')")
    scan_id = c.lastrowid
    conn.commit()
    conn.close()
    return scan_id


def log_scan_finish(scan_id, ok_count, fail_count):
    conn = get_conn()
    conn.execute("""
        UPDATE scan_log SET finished_at=datetime('now'), status='done',
        switches_ok=?, switches_fail=? WHERE id=?
    """, (ok_count, fail_count, scan_id))
    conn.commit()
    conn.close()


def get_last_scan():
    conn = get_conn()
    row = conn.execute(
        "SELECT * FROM scan_log ORDER BY id DESC LIMIT 1"
    ).fetchone()
    conn.close()
    return dict(row) if row else None


def deduplicate_links():
    """Remove duplicate bidirectional links. Always run after a scan."""
    conn = get_conn()
    before = conn.execute("SELECT COUNT(*) FROM links").fetchone()[0]
    conn.execute("""
        DELETE FROM links WHERE id NOT IN (
            SELECT MIN(id) FROM links
            GROUP BY
                CASE WHEN chassis_a < chassis_b THEN chassis_a ELSE chassis_b END,
                CASE WHEN chassis_a < chassis_b THEN chassis_b ELSE chassis_a END,
                CASE WHEN chassis_a < chassis_b THEN port_a ELSE port_b END,
                CASE WHEN chassis_a < chassis_b THEN port_b ELSE port_a END
        )
    """)
    conn.execute("DELETE FROM links WHERE chassis_a = chassis_b")
    conn.commit()
    after = conn.execute("SELECT COUNT(*) FROM links").fetchone()[0]
    conn.close()
    logger.info(f"deduplicate_links: {before} -> {after} links")


def save_node_positions(positions):
    """Save node positions from Cytoscape. positions = list of {chassis_id, x, y}"""
    conn = get_conn()
    conn.execute("DELETE FROM node_positions")
    for p in positions:
        conn.execute(
            "INSERT INTO node_positions (chassis_id, x, y) VALUES (?, ?, ?)",
            (p['chassis_id'], float(p['x']), float(p['y']))
        )
    conn.commit()
    conn.close()
    logging.getLogger(__name__).info(f"Saved {len(positions)} node positions")


def get_node_positions():
    """Returns dict of chassis_id -> (x, y), empty if none saved."""
    conn = get_conn()
    rows = conn.execute("SELECT chassis_id, x, y FROM node_positions").fetchall()
    conn.close()
    return {r["chassis_id"]: (r["x"], r["y"]) for r in rows}


def clear_node_positions():
    conn = get_conn()
    conn.execute("DELETE FROM node_positions")
    conn.commit()
    conn.close()
    logging.getLogger(__name__).info("Node positions cleared")