468d4dd6fd
C1 — Persistent connection pool: - Added CISCO_SWITCH_PERSISTENT_CONNECTION setting (default false) - Added get_connection() / close_connection() with thread-safe locking - run_command(), run_config_commands(), and test_connection() all honour persistent mode - Added cisco_switch_reconnect() tool to force a fresh connection C2 — TextFSM structured output: - Added run_command_structured() that tries TextFSM, falls back to raw JSON-wrapped text - Updated 6 diagnostic tools to return structured JSON: show_version, show_interfaces_status, show_mac_address_table, show_vlan_brief, show_ip_route, show_cdp_neighbors - Added cisco_switch_test_connection_structured() tool - Added cisco_switch_get_mcp_config() alias - Added .gitignore for __pycache__ and build artifacts
255 lines
8.1 KiB
Python
255 lines
8.1 KiB
Python
"""SSH client for Cisco IOS-XE switches via Netmiko."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import threading
|
|
from contextlib import contextmanager
|
|
from typing import Iterator
|
|
|
|
from netmiko import ConnectHandler
|
|
from netmiko.exceptions import NetmikoAuthenticationException, NetmikoTimeoutException
|
|
|
|
from cisco_switch_mcp.config import settings
|
|
from cisco_switch_mcp.guards import GuardError, require_credentials
|
|
|
|
|
|
class CiscoSwitchError(Exception):
|
|
def __init__(self, message: str, *, cause: Exception | None = None):
|
|
super().__init__(message)
|
|
self.cause = cause
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# C1 — Persistent connection pool
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_connection: ConnectHandler | None = None
|
|
_connection_lock = threading.Lock()
|
|
|
|
|
|
def get_connection() -> ConnectHandler:
|
|
"""Return a cached persistent SSH connection (thread-safe).
|
|
|
|
Creates a new connection on first call and reuses it thereafter.
|
|
Call close_connection() to tear it down.
|
|
"""
|
|
global _connection
|
|
with _connection_lock:
|
|
if _connection is None:
|
|
_connection = ConnectHandler(**_device_params())
|
|
_connection.enable()
|
|
return _connection
|
|
|
|
|
|
def close_connection() -> None:
|
|
"""Disconnect and clear the persistent SSH connection."""
|
|
global _connection
|
|
with _connection_lock:
|
|
if _connection is not None:
|
|
try:
|
|
_connection.disconnect()
|
|
except Exception:
|
|
pass
|
|
_connection = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _device_params() -> dict:
|
|
require_credentials()
|
|
secret = settings.enable_password or settings.secret or settings.password
|
|
params = {
|
|
"device_type": settings.device_type,
|
|
"host": settings.host,
|
|
"port": settings.port,
|
|
"username": settings.username,
|
|
"password": settings.password,
|
|
"secret": secret,
|
|
"timeout": settings.timeout,
|
|
"conn_timeout": settings.timeout,
|
|
"fast_cli": False,
|
|
}
|
|
if settings.session_log_path:
|
|
params["session_log"] = settings.session_log_path
|
|
return params
|
|
|
|
|
|
def _trim_output(output: str) -> str:
|
|
"""Enforce max_output_lines on raw text output."""
|
|
lines = output.splitlines()
|
|
if len(lines) > settings.max_output_lines:
|
|
trimmed = lines[: settings.max_output_lines]
|
|
trimmed.append(
|
|
f"... truncated ({len(lines) - settings.max_output_lines} more lines)"
|
|
)
|
|
return "\n".join(trimmed)
|
|
return output
|
|
|
|
|
|
@contextmanager
|
|
def connect() -> Iterator:
|
|
"""Per-call SSH connection context manager.
|
|
|
|
When CISCO_SWITCH_PERSISTENT_CONNECTION is enabled the persistent
|
|
pool is preferred; this context manager is still available for
|
|
callers that explicitly want a fresh connection.
|
|
"""
|
|
conn = None
|
|
try:
|
|
conn = ConnectHandler(**_device_params())
|
|
try:
|
|
conn.enable()
|
|
except Exception:
|
|
pass
|
|
yield conn
|
|
except NetmikoAuthenticationException as exc:
|
|
raise CiscoSwitchError("SSH authentication failed", cause=exc) from exc
|
|
except NetmikoTimeoutException as exc:
|
|
raise CiscoSwitchError(
|
|
f"SSH connection timed out to {settings.host}:{settings.port}", cause=exc
|
|
) from exc
|
|
except Exception as exc:
|
|
raise CiscoSwitchError(f"SSH connection failed: {exc}", cause=exc) from exc
|
|
finally:
|
|
if conn is not None:
|
|
try:
|
|
conn.disconnect()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _resolve_connection():
|
|
"""Return (conn, is_persistent) — the active connection."""
|
|
if settings.persistent_connection:
|
|
return get_connection(), True
|
|
# Non-persistent: we need a context manager, so use connect() as before
|
|
raise _UsePerCall # signal that callers must use the context manager
|
|
|
|
|
|
class _UsePerCall(Exception):
|
|
"""Internal sentinel used when persistent mode is off."""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public API
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def run_command(command: str, *, use_textfsm: bool = False) -> str:
|
|
"""Execute a privileged exec-mode show command."""
|
|
cmd = command.strip()
|
|
if not cmd:
|
|
raise GuardError("command is empty")
|
|
|
|
if settings.persistent_connection:
|
|
conn = get_connection()
|
|
if use_textfsm:
|
|
try:
|
|
result = conn.send_command(cmd, use_textfsm=True)
|
|
if isinstance(result, list):
|
|
return json.dumps(result, indent=2, default=str)
|
|
except Exception:
|
|
pass
|
|
output = conn.send_command(cmd, read_timeout=settings.timeout)
|
|
return _trim_output(output)
|
|
|
|
# Per-call connection (original behaviour)
|
|
with connect() as conn:
|
|
if use_textfsm:
|
|
try:
|
|
result = conn.send_command(cmd, use_textfsm=True)
|
|
if isinstance(result, list):
|
|
return json.dumps(result, indent=2, default=str)
|
|
except Exception:
|
|
pass
|
|
output = conn.send_command(cmd, read_timeout=settings.timeout)
|
|
return _trim_output(output)
|
|
|
|
|
|
def run_command_structured(command: str) -> str:
|
|
"""Execute a show command and return structured JSON via TextFSM.
|
|
|
|
Tries TextFSM first; falls back to raw text on failure.
|
|
Always returns a JSON string (structured list or {"raw_text": "…"}).
|
|
"""
|
|
try:
|
|
result = run_command(command, use_textfsm=True)
|
|
# run_command with use_textfsm=True already returns JSON string on success
|
|
return result
|
|
except Exception:
|
|
# Fall back to raw text wrapped in JSON for consistency
|
|
try:
|
|
raw = run_command(command, use_textfsm=False)
|
|
return json.dumps({"raw_text": raw}, indent=2)
|
|
except Exception as exc:
|
|
return json.dumps({"error": str(exc)}, indent=2)
|
|
|
|
|
|
def run_show_command(command: str) -> str:
|
|
"""Execute a show command with error wrapping.
|
|
|
|
Wraps run_command in a try/except and returns a JSON error
|
|
payload on failure instead of raising an exception, so that
|
|
MCP tools always return a string even when the switch is
|
|
unreachable or authentication fails.
|
|
"""
|
|
try:
|
|
return run_command(command)
|
|
except Exception as exc:
|
|
return json.dumps({"error": str(exc)}, indent=2)
|
|
|
|
|
|
def run_config_commands(commands: list[str]) -> str:
|
|
"""Apply configuration commands in config mode."""
|
|
from cisco_switch_mcp.guards import validate_config_commands
|
|
|
|
validate_config_commands(commands)
|
|
cleaned = [c.strip() for c in commands if c.strip()]
|
|
if not cleaned:
|
|
raise GuardError("no configuration commands provided")
|
|
|
|
if settings.persistent_connection:
|
|
conn = get_connection()
|
|
return conn.send_config_set(cleaned)
|
|
|
|
with connect() as conn:
|
|
output = conn.send_config_set(cleaned)
|
|
return output
|
|
|
|
|
|
def test_connection() -> dict:
|
|
require_credentials()
|
|
if settings.persistent_connection:
|
|
conn = get_connection()
|
|
else:
|
|
with connect() as conn:
|
|
hostname = conn.find_prompt()
|
|
version = conn.send_command(
|
|
"show version | include Cisco IOS XE|Model Number|System serial|uptime"
|
|
)
|
|
return {
|
|
"status": "ok",
|
|
"host": settings.host,
|
|
"prompt": hostname,
|
|
"version_summary": version.strip(),
|
|
"device_type": settings.device_type,
|
|
"persistent": False,
|
|
}
|
|
# Persistent path
|
|
hostname = conn.find_prompt()
|
|
version = conn.send_command(
|
|
"show version | include Cisco IOS XE|Model Number|System serial|uptime"
|
|
)
|
|
return {
|
|
"status": "ok",
|
|
"host": settings.host,
|
|
"prompt": hostname,
|
|
"version_summary": version.strip(),
|
|
"device_type": settings.device_type,
|
|
"persistent": True,
|
|
}
|