diff --git a/scanner.py b/scanner.py index 39eb83a..ef37f8d 100644 --- a/scanner.py +++ b/scanner.py @@ -1,6 +1,6 @@ # scanner.py - Orchestrates the full scan pipeline import logging -from nocodb_client import get_switch_ips +from nocodb_client import get_switches from ssh_client import scan_all_switches from db import ( upsert_switch, upsert_link, clear_links, @@ -39,7 +39,7 @@ def run_scan(dept: str = None, workers: int = 5, login_delay: int = 3): return # Fetch switch list from NocoDB (or fallback) - switches = get_switch_ips(dept=dept) + switches = get_switches(dept=dept) if not switches: logger.error("NocoDB returned no switches — aborting scan.") diff --git a/ssh_client.py b/ssh_client.py index e832570..778264e 100644 --- a/ssh_client.py +++ b/ssh_client.py @@ -26,9 +26,18 @@ _pt.Transport.auth_password = _auth_password_no_fallback _login_lock = threading.Semaphore(1) -def connect_and_query(ip, login_delay=3): +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, + "device_type": device_type or DEVICE_TYPE, "host": ip, "username": SSH_USERNAME, "password": SSH_PASSWORD, @@ -194,14 +203,25 @@ def _aruba_firmware(version_output): # ── Scan orchestration ──────────────────────────────────────────────────────── -def scan_all_switches(ip_list, progress_callback=None, max_workers=5, login_delay=3): +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(ip_list) + total = len(switch_list) done = 0 - _scan = partial(connect_and_query, login_delay=login_delay) + 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 = {executor.submit(_scan, ip): ip for ip in ip_list} + 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]