lldp-mapper/ssh_client.py

# ssh_client.py - SSH connections to FS, HP/Aruba ProCurve, and Dell switches
import re
import logging
import threading
import time
import socket
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import partial
from netmiko import ConnectHandler, NetmikoTimeoutException, NetmikoAuthenticationException
from config import SSH_USERNAME, SSH_PASSWORD, SSH_PORT, SSH_TIMEOUT, DEVICE_TYPE
from parser import (parse_lldp_neighbors, parse_mgmt_ip_from_interfaces,
                    parse_aruba_procurve_local, parse_aruba_procurve_neighbors,
                    normalize_mac)

logger = logging.getLogger(__name__)

# Disable paramiko's keyboard-interactive fallback so each switch gets exactly
# one auth attempt against RADIUS/AD instead of two.
import paramiko.transport as _pt
_orig_auth_password = _pt.Transport.auth_password
def _auth_password_no_fallback(self, username, password, event=None, fallback=False):
    return _orig_auth_password(self, username, password, event=event, fallback=fallback)
_pt.Transport.auth_password = _auth_password_no_fallback

# Serialise SSH logins — only one handshake/auth at a time to avoid RADIUS lockout
_login_lock = threading.Semaphore(1)


def _device_type_for(manufacturer: str) -> str:
    m = (manufacturer or "").lower()
    if any(k in m for k in ("aruba", "hp", "hewlett", "procurve")):
        return "hp_procurve"
    if "dell" in m:
        return "dell_os10"
    return DEVICE_TYPE


def connect_and_query(ip, login_delay=3, device_type=None):
    device = {
        "device_type":        device_type or DEVICE_TYPE,
        "host":               ip,
        "username":           SSH_USERNAME,
        "password":           SSH_PASSWORD,
        "port":               SSH_PORT,
        "timeout":            SSH_TIMEOUT,
        "global_delay_factor": 2,
        "fast_cli":           False,
        "conn_timeout":       SSH_TIMEOUT,
    }

    try:
        # TCP pre-flight — bail before sending any credentials if port 22 is unreachable
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(3)
            if s.connect_ex((ip, SSH_PORT)) != 0:
                logger.warning(f"{ip}: port {SSH_PORT} unreachable, skipping")
                return {"success": False, "ip": ip, "error": "Host unreachable"}

        logger.info(f"Connecting to {ip}...")

        # Allow legacy ssh-rsa keys used by FS switches
        _orig_preferred_keys = _pt.Transport._preferred_keys
        _pt.Transport._preferred_keys = (
            "ssh-rsa", "rsa-sha2-256", "rsa-sha2-512",
            "ecdsa-sha2-nistp256", "ecdsa-sha2-nistp384",
            "ecdsa-sha2-nistp521", "ssh-ed25519",
        )

        with _login_lock:
            conn = ConnectHandler(**device)
            _pt.Transport._preferred_keys = _orig_preferred_keys
            if login_delay > 0:
                time.sleep(login_delay)

        try:
            hostname = conn.find_prompt().replace('#', '').replace('>', '').strip()
            version_output = conn.send_command("show version", read_timeout=30)
            vendor = _detect_vendor(version_output)

            if vendor == 'aruba_procurve':
                chassis_id, mgmt_ip, model, firmware, neighbors = _query_aruba(conn, version_output, ip)
            else:
                chassis_id, mgmt_ip, model, firmware, neighbors = _query_fs(conn, version_output, ip)
        finally:
            conn.disconnect()

        logger.info(f"  {hostname} ({ip}) [{vendor}]: {len(neighbors)} neighbors")

        return {
            "success":     True,
            "ip":          ip,
            "hostname":    hostname,
            "chassis_id":  chassis_id,
            "mgmt_ip":     mgmt_ip,
            "description": model,
            "firmware":    firmware,
            "vendor":      vendor,
            "neighbors":   neighbors,
        }

    except NetmikoAuthenticationException:
        logger.error(f"Auth failed for {ip}")
        return {"success": False, "ip": ip, "error": "Authentication failed"}
    except NetmikoTimeoutException:
        logger.error(f"Timeout connecting to {ip}")
        return {"success": False, "ip": ip, "error": "Connection timed out"}
    except Exception as e:
        logger.error(f"Error on {ip}: {e}")
        return {"success": False, "ip": ip, "error": str(e)}


def _query_fs(conn, version_output, ip):
    """Run FS/IES switch-specific commands and return (chassis_id, mgmt_ip, model, firmware, neighbors)."""
    intf_output       = conn.send_command("show ip interface brief",    read_timeout=30)
    local_info_output = conn.send_command("show lldp local-information", read_timeout=30)
    lldp_output       = conn.send_command("show lldp neighbors",         read_timeout=30)

    mgmt_ip    = parse_mgmt_ip_from_interfaces(intf_output) or ip
    chassis_id = (_extract_local_chassis(local_info_output)
                  or _extract_mac_from_version(version_output)
                  or ip)
    chassis_id = normalize_mac(chassis_id) if '-' in chassis_id or ':' in chassis_id else chassis_id

    model    = _fs_model(version_output) or _extract_system_desc(local_info_output)
    firmware = _fs_firmware(version_output)
    neighbors = parse_lldp_neighbors(lldp_output, chassis_id, '', mgmt_ip)

    return chassis_id, mgmt_ip, model, firmware, neighbors


def _query_aruba(conn, version_output, ip):
    """Run HP/Aruba ProCurve-specific commands and return (chassis_id, mgmt_ip, model, firmware, neighbors)."""
    local_output  = conn.send_command("show lldp info local-device",  read_timeout=30)
    remote_output = conn.send_command("show lldp info remote-device", read_timeout=30)

    local_info = parse_aruba_procurve_local(local_output)
    chassis_id = local_info['chassis_id'] or _extract_mac_from_version(version_output) or ip
    mgmt_ip    = local_info['mgmt_ip'] or ip
    model      = _aruba_model(local_info.get('system_desc', ''))
    firmware   = _aruba_firmware(version_output)
    neighbors  = parse_aruba_procurve_neighbors(remote_output)

    return chassis_id, mgmt_ip, model, firmware, neighbors


# ── Vendor detection ──────────────────────────────────────────────────────────

def _detect_vendor(version_output):
    if 'Dell EMC Networking OS10' in version_output or 'Dell OS10' in version_output:
        return 'dell_os10'
    if '/ws/swbuildm/' in version_output:
        return 'aruba_procurve'
    return 'fs'


# ── FS/IES extractors ─────────────────────────────────────────────────────────

def _extract_local_chassis(output):
    m = re.search(r'Chassis ID\s*:\s*(.+)', output)
    return m.group(1).strip() if m else None


def _extract_mac_from_version(output):
    m = re.search(r'MAC Address\s*:\s*([0-9A-Fa-f]{2}(?:[:\-][0-9A-Fa-f]{2}){5})', output, re.IGNORECASE)
    return m.group(1).strip() if m else None


def _fs_model(output):
    m = re.search(r'(?:Model Name|Product)\s*:\s*([\w\-]+)', output, re.IGNORECASE)
    return f"FS {m.group(1)}" if m else ''


def _fs_firmware(output):
    fw_m   = re.search(r'Software Version\s*:\s*(\S+)', output, re.IGNORECASE)
    date_m = re.search(r'Build Date\s*:\s*(\d{2})-(\d{2})-(\d{4})T', output, re.IGNORECASE)
    if not fw_m:
        return ''
    version = fw_m.group(1)
    if date_m:
        day, month, year = date_m.group(1), date_m.group(2), date_m.group(3)
        return f"{version} ({year}-{month}-{day})"
    return version


def _extract_system_desc(output):
    m = re.search(r'System Description\s*:\s*(.*?)(?=\n\S|\Z)', output, re.DOTALL | re.IGNORECASE)
    return ' '.join(m.group(1).split()) if m else ''


# ── HP/Aruba extractors ───────────────────────────────────────────────────────

def _aruba_model(system_desc):
    """Extract model from LLDP system description, e.g. 'Aruba JL258A 2930F-8G-PoE+-2SFP+ Switch'."""
    m = re.match(r'((?:Aruba|HP)\s+\S+\s+[\w\-\+]+)', system_desc, re.IGNORECASE)
    return m.group(1) if m else system_desc.split(',')[0] if system_desc else ''


def _aruba_firmware(version_output):
    """Extract firmware version from image stamp, e.g. 'WC.16.10.0012'."""
    m = re.search(r'^\s+([A-Z]+\.\d+\.\d+\.\d+)\s*$', version_output, re.MULTILINE)
    return m.group(1).strip() if m else ''


# ── Scan orchestration ────────────────────────────────────────────────────────

def scan_all_switches(switch_list, progress_callback=None, max_workers=5, login_delay=3):
    """switch_list: list of IP strings or dicts with 'ip' and optional 'manufacturer'."""
    results = []
    total = len(switch_list)
    done  = 0

    def _submit(entry):
        if isinstance(entry, dict):
            ip = entry["ip"]
            dt = _device_type_for(entry.get("manufacturer", ""))
        else:
            ip, dt = entry, None
        return ip, executor.submit(connect_and_query, ip, login_delay=login_delay, device_type=dt)

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_ip = {}
        for entry in switch_list:
            ip, fut = _submit(entry)
            future_to_ip[fut] = ip

        for future in as_completed(future_to_ip):
            ip = future_to_ip[future]
            try:
                result = future.result()
            except Exception as e:
                result = {"success": False, "ip": ip, "error": str(e)}
            time.sleep(2)  # avoid RADIUS lockout between result processing

            results.append(result)
            done += 1

            if progress_callback:
                progress_callback(done, total, ip, result)

    return results