netx/netx_api/ne_netmiko.py

"""Shared Netmiko helpers for managed NE connect/collect."""

from __future__ import annotations

import re
import time
from typing import Any


def normalize_netmiko_device_type(device_type: str, protocol: str) -> str:
    dt = str(device_type or "").strip()
    proto = str(protocol or "ssh").strip().lower()
    low = dt.lower()
    # Raw / SecureCRT-style sessions (WebCRT quick-connect stores device_type=generic).
    if low in ("generic", "generic_ssh", "generic_telnet", "terminal_server", "generic_termserver"):
        return "generic_telnet" if proto == "telnet" or "telnet" in low else "generic_termserver_ssh"
    # Netmiko ships linux / linux_ssh but not linux_telnet — use generic_telnet.
    if low in ("linux", "linux_ssh", "linux_telnet") or low.startswith("linux_"):
        if proto == "telnet" or "telnet" in low:
            return "generic_telnet"
        return "linux_ssh"
    if "zte" in low:
        if dt == "zte":
            return f"zte_zxros_{proto}"
        if "telnet" not in low and "ssh" not in low:
            return f"{dt}_{proto}"
        return dt
    if "telnet" not in low and "ssh" not in low:
        return f"{dt}_{proto}"
    return dt


def is_cisco_ios_device_type(device_type: str) -> bool:
    dt = str(device_type or "").strip().lower()
    return "cisco_ios" in dt or dt in {"cisco_ios", "cisco_ios_ssh", "cisco_ios_telnet"}


def is_zte_device_type(device_type: str) -> bool:
    return "zte" in str(device_type or "").strip().lower()


def is_huawei_device_type(device_type: str) -> bool:
    return "huawei" in str(device_type or "").strip().lower()


def paging_disable_commands(*, vendor: str = "", device_type: str = "") -> list[str]:
    """Vendor CLI commands to disable --More-- paging (collection / LLDP / sync)."""
    dt = str(device_type or "").strip().lower()
    v = str(vendor or "").strip().lower()
    blob = f"{dt} {v}"
    # RouterOS has no pager; Cisco's command leaves an extra error/prompt redraw.
    if "mikrotik" in blob or "routeros" in blob or "switchos" in blob:
        return []
    if "huawei" in blob or "华为" in v:
        return ["screen-length 0 temporary"]
    if "hp_comware" in blob or "h3c" in blob or "comware" in blob:
        return ["screen-length disable"]
    if "juniper" in blob or "junos" in blob:
        return ["set cli screen-length 0"]
    if "nokia" in blob or "alcatel_sros" in blob or "sros" in blob:
        return ["environment no more"]
    if "alcatel_aos" in blob:
        return ["terminal length 0"]
    # Cisco / ZTE / generic SSH CLIs
    return ["terminal length 0"]


def disable_target_paging(
    conn: Any,
    *,
    vendor: str = "",
    device_type: str = "",
) -> str:
    """Send paging-off on the *current* CLI (critical after CLI hop / bastion jump).

    Nested stelnet/telnet lands on the target without Netmiko ``session_preparation``,
    so --More-- stays enabled and long ``display/show lldp`` / config dumps time out.
    Uses timing reads so a wrong hop ``base_prompt`` does not block.
    """
    cmds = paging_disable_commands(vendor=vendor, device_type=device_type)
    chunks: list[str] = []
    for cmd in cmds:
        try:
            out = conn.send_command_timing(cmd, read_timeout=15)
            chunks.append(str(out or ""))
            continue
        except Exception:
            pass
        try:
            ret = getattr(conn, "RETURN", None) or "\n"
            conn.write_channel(str(cmd) + ret)
            time.sleep(0.35)
            if hasattr(conn, "read_channel"):
                part = conn.read_channel()
                if part:
                    chunks.append(str(part))
        except Exception:
            pass
    return "".join(chunks)


def drain_read_channel(
    conn: Any,
    *,
    idle_reads: int = 4,
    pause_sec: float = 0.05,
) -> str:
    """Read until the SSH/Telnet channel is quiet.

    Cisco IOSv (and similar) often leave an extra ``R2#`` in the buffer after
    ``find_prompt`` / ``terminal length 0``. Netmiko ``send_command`` then matches
    that leftover prompt immediately and returns empty after strip — while a
    human terminal still echoes show output normally.
    """
    chunks: list[str] = []
    idle = 0
    read = getattr(conn, "read_channel", None)
    if not callable(read):
        clear = getattr(conn, "clear_buffer", None)
        if callable(clear):
            try:
                clear()
            except Exception:
                pass
        return ""
    while idle < max(1, int(idle_reads)):
        try:
            part = read()
        except Exception:
            break
        if part:
            chunks.append(str(part))
            idle = 0
            continue
        idle += 1
        time.sleep(max(0.0, float(pause_sec)))
    return "".join(chunks)


def _send_command_expect_prompt(conn: Any, cmd: str, *, read_timeout: int) -> str:
    device_type = getattr(conn, "device_type", "")
    if isinstance(device_type, str) and is_cisco_ios_device_type(device_type):
        # IOSv has the same late-prompt race as RouterOS: auto_find_prompt sends
        # a newline after our drain and its reply can terminate the next show.
        # Use the session's cached hostname, but require a full prompt line so
        # "R2 uptime ..." inside show version cannot end the command early.
        #
        # cmd_verify must stay False: Netmiko hardcodes a 10s echo wait, and
        # IOSv often echoes ``show version | include ...`` slower than that
        # (live: ~13s) even when the command itself succeeds. Correctness comes
        # from drain_read_channel + full-line expect_string, not echo verify.
        base_prompt = getattr(conn, "base_prompt", "")
        if not isinstance(base_prompt, str) or not base_prompt.strip():
            raise ValueError("cisco_base_prompt_missing")
        prompt_pattern = rf"(?m:^{re.escape(base_prompt.strip())}[>#][ \t]*$)"
        out = str(
            conn.send_command(
                command_string=cmd,
                read_timeout=read_timeout,
                auto_find_prompt=False,
                expect_string=prompt_pattern,
                cmd_verify=False,
            )
            or ""
        )
        # Leftover-prompt races still yield a silent empty body with ok=true at
        # higher layers. Fail loud here — never re-send the command.
        if not out.strip():
            raise ValueError("cisco_empty_show_output")
        return out
    if isinstance(device_type, str) and device_type.lower().startswith("mikrotik"):
        # find_prompt sends an extra Enter *after* our channel drain. RouterOS
        # can repaint that old prompt after send_command starts, ending the read
        # before the command's output arrives. Use its complete prompt directly.
        # Keep echo verification off: long RouterOS scripts wrap at PTY width.
        prompt_pattern = r"(?m)^\[[^\r\n\]]+\](?:[ \t]+/[^\r\n>]*)?[ \t]*>[ \t]*$"
        return str(
            conn.send_command(
                command_string=cmd,
                read_timeout=read_timeout,
                auto_find_prompt=False,
                expect_string=prompt_pattern,
                cmd_verify=False,
            )
            or ""
        )
    try:
        return str(
            conn.send_command(
                command_string=cmd,
                read_timeout=read_timeout,
                cmd_verify=False,
            )
            or ""
        )
    except TypeError:
        # Older Netmiko without cmd_verify kwarg.
        return str(conn.send_command(command_string=cmd, read_timeout=read_timeout) or "")


def send_show_command(conn: Any, command: str, *, read_timeout: int = 120) -> str:
    """Send one CLI command via ``send_command`` (wait for device prompt).

    Never re-sends the same command on empty/short output: this path is shared by
    show collection and open-policy config/scripts (MikroTik ``add``, etc.). A
    leftover-prompt empty read must be fixed by draining *before* send, not by
    retrying — retries duplicate side-effecting commands.

    Prefer ``send_command`` over ``send_command_timing`` for long config dumps:
    idle during ``Building configuration...`` truncates timing reads.

    Cisco uses ``cmd_verify=False`` with a full-line prompt expect after drain;
    Netmiko's echo verify is hardcoded to 10s and false-fails slow IOSv echoes.
    Other drivers also use ``cmd_verify=False``: the default echo check often
    raises ``Pattern not detected`` on hop / slow echo paths.
    """
    cmd = str(command or "").strip()
    if not cmd:
        return ""

    # IOSv / hop: find_prompt / paging-off often leave an extra prompt in the
    # channel; drain so send_command does not match that leftover immediately.
    device_type = getattr(conn, "device_type", "")
    if isinstance(device_type, str) and is_cisco_ios_device_type(device_type):
        drain_read_channel(conn, idle_reads=6, pause_sec=0.08)
    else:
        drain_read_channel(conn)
    return _send_command_expect_prompt(conn, cmd, read_timeout=read_timeout)