netx/netx_api/ne_connect.py

from __future__ import annotations

import logging
import re
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from typing import Any

from netmiko import ConnectHandler

from .config import settings
from .db import SessionLocal
from .models import ManagedNE
from .ne_crypto import CredentialCryptoError, decrypt_secret
from .ne_netmiko import normalize_netmiko_device_type
from .ne_service import get_device_credentials

_log = logging.getLogger("netx.ne.connect")
_executor: ThreadPoolExecutor | None = None


def _executor_pool() -> ThreadPoolExecutor:
    global _executor
    if _executor is None:
        workers = max(1, int(settings.ne_connect_max_workers or 5))
        _executor = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="ne-connect")
    return _executor


def hostname_probe_command(device_type: str, vendor: str) -> str | None:
    """
    Per-vendor CLI to read system name (ported from legacy connect.extract_dev_command).
    ZTE: rely on login prompt / empty command path.
    """
    dt = str(device_type or "").lower()
    v = str(vendor or "").lower()
    if "huawei" in dt or v == "huawei":
        return "display current-configuration | include sysname"
    if "juniper" in dt or v == "juniper":
        return "show system host-name"
    if "cisco" in dt or v == "cisco":
        return "show hostname"
    return None


def parse_hostname_from_output(
    device_type: str,
    vendor: str,
    output: str,
    prompt: str = "",
) -> str | None:
    """
    Parse device name from command output or prompt (legacy connect.extract_hostname).
    """
    dt = str(device_type or "").lower()
    v = str(vendor or "").lower()
    text = str(output or "")

    if "huawei" in dt or v == "huawei":
        m = re.search(r"sysname\s+(\S+)", text, re.IGNORECASE)
        if m:
            return m.group(1).strip()

    if "juniper" in dt or v == "juniper":
        m = re.search(r"host-name\s+(\S+)", text, re.IGNORECASE)
        if m:
            return m.group(1).strip().rstrip(";")
        m = re.search(r"^\s*name\s+(\S+)", text, re.IGNORECASE | re.MULTILINE)
        if m:
            return m.group(1).strip().rstrip(";")

    if "cisco" in dt or v == "cisco":
        lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
        for ln in reversed(lines):
            if ln.startswith("%") or "invalid" in ln.lower():
                continue
            token = ln.split()[0].strip("<>[]")
            if token:
                return token

    if "zte" in dt or v == "zte":
        lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
        if lines:
            last = lines[-1].strip()
            if last and len(last) <= 128 and not last.startswith("%"):
                return last

    cleaned = _clean_prompt_hostname(prompt)
    if cleaned:
        return cleaned
    return None


def _clean_prompt_hostname(prompt: str) -> str | None:
    p = str(prompt or "").strip()
    if not p:
        return None
    p = re.sub(r"[\s#>$]+\s*$", "", p).strip()
    p = re.sub(r"^[<\[]|[>\]]$", "", p).strip()
    if not p or p.lower() in (">", "#"):
        return None
    return p[:256]


def _probe_device(creds: dict[str, Any]) -> tuple[str, str, str | None]:
    """Login via Netmiko, probe hostname, return (status, message, discovered_name)."""
    device_type = normalize_netmiko_device_type(creds["device_type"], creds["protocol"])
    vendor = str(creds.get("vendor") or "")
    dev: dict[str, Any] = {
        "device_type": device_type,
        "host": creds["ip_address"],
        "username": creds["username"],
        "password": creds["password"],
        "port": int(creds["port"] or 22),
        "conn_timeout": int(settings.ne_connect_timeout_sec or 30),
        "auth_timeout": int(settings.ne_connect_timeout_sec or 30),
        "banner_timeout": int(settings.ne_connect_timeout_sec or 30),
    }
    secret = str(creds.get("enable_secret") or "").strip()
    if secret:
        dev["secret"] = secret
    try:
        with ConnectHandler(**dev) as conn:
            prompt = str(conn.find_prompt() or "")
            command = hostname_probe_command(creds["device_type"], vendor)
            output = ""
            if command:
                output = conn.send_command(command_string=command, read_timeout=30)
            hostname = parse_hostname_from_output(creds["device_type"], vendor, output, prompt)
            if hostname:
                return "pass", f"connected: {hostname}", hostname
            if command:
                return "pass", "connected (hostname not parsed)", None
            fallback = _clean_prompt_hostname(prompt)
            if fallback:
                return "pass", f"connected: {fallback}", fallback
            return "pass", "connected", None
    except Exception as exc:
        msg = str(exc).split("\n")[0][:480]
        return "fail", msg, None


def _update_row(ne_id: str, status: str, message: str, discovered_name: str | None = None) -> None:
    db = SessionLocal()
    try:
        row = db.get(ManagedNE, ne_id)
        if not row:
            return
        row.connect_status = status
        row.connect_message = str(message or "")[:500]
        row.connect_tested_at = datetime.utcnow()
        if discovered_name:
            row.name = discovered_name[:256]
        row.updated_at = datetime.utcnow()
        db.commit()
    finally:
        db.close()


def _run_single(ne_id: str) -> None:
    db = SessionLocal()
    try:
        row = db.get(ManagedNE, ne_id)
        if not row:
            return
        row.connect_status = "testing"
        row.connect_message = ""
        row.updated_at = datetime.utcnow()
        db.commit()
        try:
            creds = get_device_credentials(row)
        except CredentialCryptoError as exc:
            _update_row(ne_id, "fail", str(exc))
            return
        status, message, discovered = _probe_device(creds)
        _update_row(ne_id, status, message, discovered)
    except Exception as exc:
        _log.exception("connect test failed for %s", ne_id)
        _update_row(ne_id, "fail", str(exc)[:480])
    finally:
        db.close()


def schedule_connect_tests(ne_ids: list[str]) -> int:
    pool = _executor_pool()
    submitted = 0
    for ne_id in ne_ids:
        ne_id = str(ne_id or "").strip()
        if not ne_id:
            continue
        pool.submit(_run_single, ne_id)
        submitted += 1
    return submitted