C1: persistent SSH connection pool, C2: TextFSM structured output
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
This commit is contained in:
@@ -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,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user