From 468d4dd6fd7e4c44991a98727247728dd7e87dad Mon Sep 17 00:00:00 2001 From: root Date: Sat, 4 Jul 2026 20:41:26 +0000 Subject: [PATCH] C1: persistent SSH connection pool, C2: TextFSM structured output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .gitignore | 7 + src/cisco_switch_mcp/client.py | 149 ++++++++++++++++++++-- src/cisco_switch_mcp/config.py | 3 + src/cisco_switch_mcp/tools/diagnostics.py | 4 +- src/cisco_switch_mcp/tools/interfaces.py | 4 +- src/cisco_switch_mcp/tools/l2.py | 6 +- src/cisco_switch_mcp/tools/system.py | 56 +++++++- src/cisco_switch_mcp/tools/vlan.py | 4 +- 8 files changed, 209 insertions(+), 24 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9019739 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +*.pyc +*.pyo +*.egg-info/ +dist/ +build/ +.eggs/ diff --git a/src/cisco_switch_mcp/client.py b/src/cisco_switch_mcp/client.py index f9ad425..b4ac4df 100644 --- a/src/cisco_switch_mcp/client.py +++ b/src/cisco_switch_mcp/client.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import threading from contextlib import contextmanager from typing import Iterator @@ -19,6 +20,45 @@ class CiscoSwitchError(Exception): 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 @@ -38,8 +78,26 @@ def _device_params() -> dict: 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()) @@ -64,30 +122,71 @@ def connect() -> Iterator: 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): - import json - return json.dumps(result, indent=2, default=str) except Exception: pass output = conn.send_command(cmd, read_timeout=settings.timeout) - 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 + 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: @@ -112,6 +211,11 @@ def run_config_commands(commands: list[str]) -> str: 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 @@ -119,13 +223,32 @@ def run_config_commands(commands: list[str]) -> str: def test_connection() -> dict: require_credentials() - with connect() as conn: - hostname = conn.find_prompt() - version = conn.send_command("show version | include Cisco IOS XE|Model Number|System serial|uptime") + 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, } diff --git a/src/cisco_switch_mcp/config.py b/src/cisco_switch_mcp/config.py index e579e88..4b0a513 100644 --- a/src/cisco_switch_mcp/config.py +++ b/src/cisco_switch_mcp/config.py @@ -23,5 +23,8 @@ class Settings(BaseSettings): block_config_replace: bool = True max_output_lines: int = 500 + # C1: Enable persistent SSH connection pool (default false for backward compat) + persistent_connection: bool = False + settings = Settings() diff --git a/src/cisco_switch_mcp/tools/diagnostics.py b/src/cisco_switch_mcp/tools/diagnostics.py index 9df6175..c00e8fb 100644 --- a/src/cisco_switch_mcp/tools/diagnostics.py +++ b/src/cisco_switch_mcp/tools/diagnostics.py @@ -4,7 +4,7 @@ from __future__ import annotations from mcp.server.fastmcp import FastMCP -from cisco_switch_mcp.client import run_show_command +from cisco_switch_mcp.client import run_command_structured, run_show_command def register(mcp: FastMCP) -> None: @@ -52,7 +52,7 @@ def register(mcp: FastMCP) -> None: @mcp.tool() def cisco_switch_show_ip_route() -> str: """Get IP routing table on the switch.""" - return run_show_command("show ip route") + return run_command_structured("show ip route") @mcp.tool() def cisco_switch_show_users() -> str: diff --git a/src/cisco_switch_mcp/tools/interfaces.py b/src/cisco_switch_mcp/tools/interfaces.py index c5e8158..4cbcb05 100644 --- a/src/cisco_switch_mcp/tools/interfaces.py +++ b/src/cisco_switch_mcp/tools/interfaces.py @@ -4,7 +4,7 @@ from __future__ import annotations from mcp.server.fastmcp import FastMCP -from cisco_switch_mcp.client import run_show_command +from cisco_switch_mcp.client import run_command_structured, run_show_command from cisco_switch_mcp.guards import validate_interface @@ -12,7 +12,7 @@ def register(mcp: FastMCP) -> None: @mcp.tool() def cisco_switch_show_interfaces_status() -> str: """Get brief interface status (port, VLAN, speed, duplex, status).""" - return run_show_command("show interfaces status") + return run_command_structured("show interfaces status") @mcp.tool() def cisco_switch_show_interfaces_description() -> str: diff --git a/src/cisco_switch_mcp/tools/l2.py b/src/cisco_switch_mcp/tools/l2.py index 43cf0c4..f49b864 100644 --- a/src/cisco_switch_mcp/tools/l2.py +++ b/src/cisco_switch_mcp/tools/l2.py @@ -4,7 +4,7 @@ from __future__ import annotations from mcp.server.fastmcp import FastMCP -from cisco_switch_mcp.client import run_show_command +from cisco_switch_mcp.client import run_command_structured, run_show_command from cisco_switch_mcp.guards import validate_interface @@ -12,7 +12,7 @@ def register(mcp: FastMCP) -> None: @mcp.tool() def cisco_switch_show_mac_address_table() -> str: """Get MAC address table (CAM table).""" - return run_show_command("show mac address-table") + return run_command_structured("show mac address-table") @mcp.tool() def cisco_switch_show_mac_address_table_vlan(vlan_id: int) -> str: @@ -28,7 +28,7 @@ def register(mcp: FastMCP) -> None: @mcp.tool() def cisco_switch_show_cdp_neighbors() -> str: """Get CDP neighbor discovery table.""" - return run_show_command("show cdp neighbors detail") + return run_command_structured("show cdp neighbors detail") @mcp.tool() def cisco_switch_show_lldp_neighbors() -> str: diff --git a/src/cisco_switch_mcp/tools/system.py b/src/cisco_switch_mcp/tools/system.py index 6a53ca1..e5b8b7d 100644 --- a/src/cisco_switch_mcp/tools/system.py +++ b/src/cisco_switch_mcp/tools/system.py @@ -6,7 +6,13 @@ import json from mcp.server.fastmcp import FastMCP -from cisco_switch_mcp.client import CiscoSwitchError, run_show_command, test_connection +from cisco_switch_mcp.client import ( + CiscoSwitchError, + close_connection, + run_command_structured, + run_show_command, + test_connection, +) from cisco_switch_mcp.config import settings @@ -22,10 +28,16 @@ def register(mcp: FastMCP) -> None: "username_configured": bool(settings.username), "allow_writes": settings.allow_writes, "dry_run": settings.dry_run, + "persistent_connection": settings.persistent_connection, }, indent=2, ) + @mcp.tool() + def cisco_switch_get_mcp_config() -> str: + """Get MCP configuration summary.""" + return cisco_switch_get_config() + @mcp.tool() def cisco_switch_test_connection() -> str: """Test SSH connectivity and authentication to the switch.""" @@ -34,10 +46,50 @@ def register(mcp: FastMCP) -> None: except CiscoSwitchError as exc: return json.dumps({"status": "error", "message": str(exc)}, indent=2) + @mcp.tool() + def cisco_switch_test_connection_structured() -> str: + """Test SSH connectivity using TextFSM-structured show version output.""" + try: + result = run_command_structured("show version") + return json.dumps( + { + "status": "ok", + "host": settings.host, + "device_type": settings.device_type, + "persistent": settings.persistent_connection, + "structured_version": json.loads(result) if result else {}, + }, + indent=2, + ) + except Exception as exc: + return json.dumps({"status": "error", "message": str(exc)}, indent=2) + + @mcp.tool() + def cisco_switch_reconnect() -> str: + """Close and re-establish the persistent SSH connection (when persistent mode is on).""" + close_connection() + try: + # Force a new connection on next call + from cisco_switch_mcp.client import get_connection + + conn = get_connection() + prompt = conn.find_prompt() + return json.dumps( + { + "status": "ok", + "message": "Reconnected successfully.", + "prompt": prompt, + "persistent": True, + }, + indent=2, + ) + except Exception as exc: + return json.dumps({"status": "error", "message": str(exc)}, indent=2) + @mcp.tool() def cisco_switch_show_version() -> str: """Get IOS-XE version, model, serial number, and uptime.""" - return run_show_command("show version") + return run_command_structured("show version") @mcp.tool() def cisco_switch_show_inventory() -> str: diff --git a/src/cisco_switch_mcp/tools/vlan.py b/src/cisco_switch_mcp/tools/vlan.py index 9ed0073..869337c 100644 --- a/src/cisco_switch_mcp/tools/vlan.py +++ b/src/cisco_switch_mcp/tools/vlan.py @@ -4,14 +4,14 @@ from __future__ import annotations from mcp.server.fastmcp import FastMCP -from cisco_switch_mcp.client import run_show_command +from cisco_switch_mcp.client import run_command_structured, run_show_command def register(mcp: FastMCP) -> None: @mcp.tool() def cisco_switch_show_vlan_brief() -> str: """List VLANs and port assignments (show vlan brief).""" - return run_show_command("show vlan brief") + return run_command_structured("show vlan brief") @mcp.tool() def cisco_switch_show_vlan(id: int = 0) -> str: