from mcp.server.fastmcp import FastMCP from typing import Dict, List, Optional, Any import asyncio import logging import sys import os import glob import json # 导入我们创建的模块 from tool_manager import tool_manager, ToolCategory from network_devices import ( device_manager, NetworkDevice, DeviceCredential, DeviceVendor, DeviceType, DeviceStatus, ConnectionProtocol ) from device_connector import connection_manager, CommandResult, SCRAPLI_IMPORT_SUCCESS # 导入资源管理模块 from mcp_resources import ( resource_manager, ResourceType ) # 导入新的模板系统 from template_system import ( template_manager, render_template_with_resources, Message, UserMessage, AssistantMessage, SystemMessage ) # 导入设备命令提示模板系统 import device_prompts # 导入拓扑发现模块 from topology_discovery_improved import create_improved_topology_discovery # 导入网络扫描模块 from network_scanner import create_network_scanner # 引入相同的日志格式化处理 class JsonFormatter(logging.Formatter): def __init__(self, fmt=None, datefmt=None, style='%'): super().__init__(fmt, datefmt, style) def format(self, record): log_record = super().format(record) return log_record.encode('utf-8', errors='replace').decode('utf-8') # 配置日志格式 logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.StreamHandler(stream=sys.stdout) ] ) # 设置所有处理器使用UTF-8编码格式化 for handler in logging.root.handlers: if isinstance(handler, logging.StreamHandler): handler.setFormatter(JsonFormatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) logger = logging.getLogger("netbrain_mcp") # 创建一个MCP服务器 mcp = FastMCP("NetBrain MCP") # 创建拓扑发现引擎实例 topology_discovery = create_improved_topology_discovery(device_manager, connection_manager) # 创建网络扫描器实例 network_scanner = create_network_scanner(device_manager) # 工具分类：设备管理 @mcp.tool() async def list_devices( vendor: Optional[str] = None, device_type: Optional[str] = None, status: Optional[str] = None, tag: Optional[str] = None ) -> List[Dict[str, Any]]: """ 列出网络设备，支持过滤 Args: vendor: 设备厂商，可选值：cisco, huawei, h3c, juniper, arista, fortinet, checkpoint, other device_type: 设备类型，可选值：router, switch, firewall, load_balancer, wireless_controller, access_point, other status: 设备状态，可选值：online, offline, unreachable, maintenance, unknown tag: 标签过滤 Returns: 设备列表 """ # 转换过滤参数 vendor_enum = DeviceVendor(vendor) if vendor else None device_type_enum = DeviceType(device_type) if device_type else None status_enum = DeviceStatus(status) if status else None # 获取设备列表 devices = device_manager.list_devices( vendor=vendor_enum, device_type=device_type_enum, status=status_enum, tag=tag ) # 转换为字典列表 return [device.to_dict() for device in devices] @mcp.tool() async def add_device( name: str, ip_address: str, device_type: str, vendor: str, platform: str = "", model: str = "", os_version: str = "", location: str = "", description: str = "", tags: str = "" ) -> Dict[str, Any]: """ 添加新的网络设备 Args: name: 设备名称 ip_address: 设备IP地址 device_type: 设备类型，可选值：router, switch, firewall, load_balancer, wireless_controller, access_point, other vendor: 设备厂商，可选值：cisco, huawei, h3c, juniper, arista, fortinet, checkpoint, other platform: 设备平台(scrapli平台类型)，如 cisco_iosxe, huawei_vrp, juniper_junos 等 model: 设备型号 os_version: 操作系统版本 location: 设备位置 description: 设备描述 tags: 设备标签，使用逗号分隔多个标签 Returns: 新添加的设备信息 """ # 转换参数 device_type_enum = DeviceType(device_type) vendor_enum = DeviceVendor(vendor) # 如果未提供platform，根据vendor推断 if not platform: if vendor.lower() == "cisco": platform = "cisco_iosxe" elif vendor.lower() == "huawei": platform = "huawei_vrp" elif vendor.lower() == "juniper": platform = "juniper_junos" elif vendor.lower() == "arista": platform = "arista_eos" else: platform = f"{vendor.lower()}" # 处理标签 tag_list = [] if tags: tag_list = [tag.strip() for tag in tags.split(',') if tag.strip()] # 创建设备对象 device = NetworkDevice( name=name, ip_address=ip_address, device_type=device_type_enum, vendor=vendor_enum, platform=platform, model=model, os_version=os_version, location=location, description=description, tags=tag_list ) # 添加设备 device_id = device_manager.add_device(device) return device.to_dict() @mcp.tool() async def get_device(device_id: str) -> Optional[Dict[str, Any]]: """ 获取设备详细信息 Args: device_id: 设备ID Returns: 设备信息或None """ device = device_manager.get_device(device_id) if device: return device.to_dict() return None @mcp.tool() async def update_device( device_id: str, name: str = "", ip_address: str = "", model: str = "", os_version: str = "", status: str = "", location: str = "", description: str = "", tags: str = "" ) -> Optional[Dict[str, Any]]: """ 更新设备信息 Args: device_id: 设备ID name: 设备名称 ip_address: 设备IP地址 model: 设备型号 os_version: 操作系统版本 status: 设备状态，可选值：online, offline, unreachable, maintenance, unknown location: 设备位置 description: 设备描述 tags: 设备标签，使用逗号分隔多个标签 Returns: 更新后的设备信息或None """ # 准备更新参数 update_kwargs = {} if name: update_kwargs["name"] = name if ip_address: update_kwargs["ip_address"] = ip_address if model: update_kwargs["model"] = model if os_version: update_kwargs["os_version"] = os_version if status: update_kwargs["status"] = DeviceStatus(status) if location: update_kwargs["location"] = location if description: update_kwargs["description"] = description if tags: # 将逗号分隔的标签转换为列表 tag_list = [tag.strip() for tag in tags.split(',') if tag.strip()] update_kwargs["tags"] = tag_list # 更新设备 device = device_manager.update_device(device_id, **update_kwargs) if device: return device.to_dict() return None @mcp.tool() async def delete_device(device_id: str) -> bool: """ 删除设备 Args: device_id: 设备ID Returns: 是否删除成功 """ return device_manager.delete_device(device_id) # 工具分类：凭据管理 @mcp.tool() async def add_credential( name: str, username: str, password: str = "", protocol: str = "ssh", port: Optional[int] = None, enable_password: str = "", ssh_key_file: str = "" ) -> Dict[str, Any]: """ 添加设备凭据 Args: name: 凭据名称 username: 用户名 password: 密码 protocol: 连接协议，可选值：ssh, telnet, snmp, http, https, netconf port: 端口号，可选（默认SSH=22, Telnet=23） enable_password: 特权模式密码（思科设备） ssh_key_file: SSH密钥文件路径 Returns: 凭据ID """ # 确保密码或SSH密钥的要求基于协议 if protocol.lower() == "ssh" and not password and not ssh_key_file: return {"success": False, "message": "SSH连接需要提供密码或SSH密钥文件"} elif protocol.lower() != "ssh" and not password: return {"success": False, "message": f"{protocol.upper()}连接需要提供密码"} # 转换参数 protocol_enum = ConnectionProtocol(protocol.lower()) # 设置默认端口 if port is None: if protocol_enum == ConnectionProtocol.SSH: port = 22 elif protocol_enum == ConnectionProtocol.TELNET: port = 23 elif protocol_enum == ConnectionProtocol.HTTP: port = 80 elif protocol_enum == ConnectionProtocol.HTTPS: port = 443 else: port = 0 # 记录凭据信息(不包含敏感信息) logger.info(f"添加凭据: name={name}, username={username}, protocol={protocol}, port={port}") if ssh_key_file: logger.info(f"使用SSH密钥文件: {ssh_key_file}") # 创建凭据对象 credential = DeviceCredential( name=name, username=username, password=password, protocol=protocol_enum, port=port, enable_password=enable_password, ssh_key_file=ssh_key_file ) # 添加凭据 credential_id = device_manager.add_credential(credential) return {"success": True, "id": credential_id, "name": name} @mcp.tool() async def list_credentials() -> List[Dict[str, Any]]: """ 列出所有设备凭据 Returns: 凭据列表 """ credentials = device_manager.list_credentials() return [ { "id": cred.id, "name": cred.name, "username": cred.username, "password": cred.password, "protocol": cred.protocol.value, "port": cred.port, "ssh_key_file": cred.ssh_key_file } for cred in credentials ] # 工具分类：设备连接 @mcp.tool() async def connect_device(device_id: str, credential_id: str) -> Dict[str, Any]: """ 连接到网络设备 Args: device_id: 设备ID credential_id: 凭据ID Returns: 连接结果 """ # 获取设备和凭据 device = device_manager.get_device(device_id) credential = device_manager.get_credential(credential_id) if not device: return {"success": False, "message": f"设备不存在: {device_id}"} if not credential: return {"success": False, "message": f"凭据不存在: {credential_id}"} # 记录连接信息 logger.info(f"尝试连接设备: {device.name} ({device.ip_address}) 使用凭据: {credential.name}") logger.info(f"连接协议: {credential.protocol.value}, 端口: {credential.port or '默认'}") # 连接设备 success, error = await connection_manager.connect_device(device, credential) if success: # 获取更多的连接信息 prompt = "" version_info = "" try: # 获取连接的connector connection_key = f"{device_id}_{credential_id}" connector = connection_manager.active_connections.get(connection_key) if connector and connector.connection: # 获取提示符 if hasattr(connector.connection, "get_prompt"): if asyncio.iscoroutinefunction(connector.connection.get_prompt): prompt = await connector.connection.get_prompt() else: prompt = await asyncio.to_thread(connector.connection.get_prompt) logger.info(f"成功获取设备提示符: {prompt}") # 使用与test_scrapli_connection相同的方式获取版本信息 try: from scrapli import Scrapli # 确定平台类型 platform = device.platform if hasattr(device, 'platform') and device.platform else "" # 根据平台类型选择命令 if "huawei" in platform.lower(): command = "display version" elif "cisco" in platform.lower(): command = "show version" else: command = "show version" logger.info(f"尝试获取版本信息，平台: {platform}，命令: {command}") # 直接使用Scrapli连接对象发送命令 if hasattr(connector.connection, "send_command"): # 等待1秒确保连接稳定 await asyncio.sleep(1) if asyncio.iscoroutinefunction(connector.connection.send_command): resp = await connector.connection.send_command(command) if hasattr(resp, "result"): version_info = resp.result else: version_info = str(resp) else: # 同步调用需要运行在线程中 resp = await asyncio.to_thread(connector.connection.send_command, command) if hasattr(resp, "result"): version_info = resp.result else: version_info = str(resp) logger.info(f"成功获取版本信息，长度: {len(version_info)}") # 限制长度 if len(version_info) > 500: version_info = version_info[:500] + "..." else: logger.warning("连接对象没有send_command方法") version_info = "无法获取版本信息：连接对象不支持发送命令" except Exception as e: logger.warning(f"获取版本信息时出错: {str(e)}") # 回退到使用connection_manager发送命令 try: result, _ = await connection_manager.send_command( device_id=device_id, credential_id=credential_id, command=command, timeout=60 ) if result and result.success: version_info = result.output if len(version_info) > 500: version_info = version_info[:500] + "..." else: version_info = "无法获取版本信息" except Exception as e2: logger.error(f"备用方法获取版本信息也失败: {str(e2)}") version_info = "无法获取版本信息" except Exception as e: logger.warning(f"获取设备详细信息时出错: {str(e)}") version_info = "获取版本信息时出错" return { "success": True, "message": f"成功连接到设备: {device.name} ({device.ip_address})", "protocol": credential.protocol.value, "prompt": prompt, "version_info": version_info } else: logger.error(f"连接设备失败: {device.name} ({device.ip_address}), 错误: {error}") return {"success": False, "message": f"连接设备失败: {error}"} @mcp.tool() async def disconnect_device(device_id: str, credential_id: str) -> Dict[str, Any]: """ 断开与网络设备的连接 Args: device_id: 设备ID credential_id: 凭据ID Returns: 断开连接结果 """ # 断开连接 success, error = await connection_manager.disconnect_device(device_id, credential_id) if success: return {"success": True, "message": "成功断开连接"} else: return {"success": False, "message": f"断开连接失败: {error}"} @mcp.tool() async def send_command( device_id: str, credential_id: str, command: str, timeout: int = 30 ) -> Dict[str, Any]: """ 向网络设备发送命令 Args: device_id: 设备ID credential_id: 凭据ID command: 要执行的命令 timeout: 命令超时时间（秒） Returns: 命令执行结果 """ # 发送命令 result, error = await connection_manager.send_command( device_id=device_id, credential_id=credential_id, command=command, timeout=timeout ) if result: return { "success": result.success, "command": result.command, "output": result.output, "error_message": result.error_message, "execution_time": result.execution_time.isoformat() } else: return { "success": False, "command": command, "output": "", "error_message": error, "execution_time": None } @mcp.tool() async def send_commands( device_id: str, credential_id: str, commands: str, timeout: int = 30 ) -> List[Dict[str, Any]]: """ 向网络设备发送多个命令 Args: device_id: 设备ID credential_id: 凭据ID commands: 要执行的命令列表，使用分号(;)分隔多个命令 timeout: 每个命令的超时时间（秒） Returns: 命令执行结果列表 """ # 将分号分隔的命令字符串转换为列表 command_list = [cmd.strip() for cmd in commands.split(';') if cmd.strip()] if not command_list: return [{"success": False, "message": "未提供有效的命令"}] # 发送命令 results, error = await connection_manager.send_commands( device_id=device_id, credential_id=credential_id, commands=command_list, timeout=timeout ) if results: return [ { "success": result.success, "command": result.command, "output": result.output, "error_message": result.error_message, "execution_time": result.execution_time.isoformat() } for result in results ] else: return [ { "success": False, "command": command, "output": "", "error_message": error, "execution_time": None } for command in command_list ] @mcp.tool() async def get_active_connections() -> List[Dict[str, Any]]: """ 获取活动连接列表 Returns: 活动连接信息列表 """ return connection_manager.get_active_connections() # 工具分类：拓扑发现 @mcp.tool() async def discover_topology(device_ids: str) -> Dict[str, Any]: """ 从指定设备开始发现网络拓扑 Args: device_ids: 设备ID列表，用逗号分隔，如"device1,device2" Returns: 拓扑发现结果 """ try: # 解析设备ID列表 if isinstance(device_ids, str): device_id_list = [device_id.strip() for device_id in device_ids.split(',') if device_id.strip()] else: device_id_list = device_ids if not device_id_list: return {"success": False, "message": "请提供至少一个设备ID"} logger.info(f"开始拓扑发现，设备列表: {device_id_list}") # 验证设备是否存在 valid_devices = [] for device_id in device_id_list: device = device_manager.get_device(device_id) if device: valid_devices.append(device_id) else: logger.warning(f"设备不存在: {device_id}") if not valid_devices: return {"success": False, "message": "没有找到有效的设备"} # 开始拓扑发现 topology = await topology_discovery.discover_topology_from_devices(valid_devices) return { "success": True, "message": f"拓扑发现完成，发现{len(topology.nodes)}个节点，{len(topology.links)}条链路", "topology": topology.to_dict() } except Exception as e: logger.error(f"拓扑发现失败: {e}") return {"success": False, "message": f"拓扑发现失败: {str(e)}"} @mcp.tool() async def get_topology() -> Dict[str, Any]: """ 获取当前网络拓扑 Returns: 当前拓扑信息 """ try: topology = topology_discovery.get_topology() return { "success": True, "topology": topology.to_dict() } except Exception as e: logger.error(f"获取拓扑失败: {e}") return {"success": False, "message": f"获取拓扑失败: {str(e)}"} @mcp.tool() async def clear_topology() -> Dict[str, Any]: """ 清空拓扑数据 Returns: 操作结果 """ try: topology_discovery.clear_topology() return {"success": True, "message": "拓扑数据已清空"} except Exception as e: logger.error(f"清空拓扑失败: {e}") return {"success": False, "message": f"清空拓扑失败: {str(e)}"} @mcp.tool() async def get_device_neighbors(device_id: str) -> Dict[str, Any]: """ 获取指定设备的邻居设备 Args: device_id: 设备ID Returns: 邻居设备列表 """ try: topology = topology_discovery.get_topology() neighbors = topology.get_device_neighbors(device_id) # 获取邻居设备的详细信息 neighbor_details = [] for neighbor_id in neighbors: device = device_manager.get_device(neighbor_id) if device: neighbor_details.append(device.to_dict()) return { "success": True, "device_id": device_id, "neighbors": neighbor_details } except Exception as e: logger.error(f"获取设备邻居失败: {e}") return {"success": False, "message": f"获取设备邻居失败: {str(e)}"} @mcp.tool() async def discover_device_neighbors(device_id: str) -> Dict[str, Any]: """ 发现单个设备的邻居（实时发现） Args: device_id: 设备ID Returns: 发现的邻居信息 """ try: device = device_manager.get_device(device_id) if not device: return {"success": False, "message": f"设备不存在: {device_id}"} logger.info(f"开始发现设备 {device.name} 的邻居") # 发现设备邻居 neighbors, interfaces = await topology_discovery.discover_device_neighbors(device) return { "success": True, "device_id": device_id, "device_name": device.name, "neighbors": neighbors, "interfaces": [iface.to_dict() for iface in interfaces], "neighbor_count": len(neighbors), "interface_count": len(interfaces) } except Exception as e: logger.error(f"发现设备邻居失败: {e}") return {"success": False, "message": f"发现设备邻居失败: {str(e)}"} @mcp.tool() async def get_topology_statistics() -> Dict[str, Any]: """ 获取拓扑统计信息 Returns: 拓扑统计数据 """ try: topology = topology_discovery.get_topology() # 统计不同协议的链路数量 protocol_stats = {} for link in topology.links: protocol = link.protocol.value protocol_stats[protocol] = protocol_stats.get(protocol, 0) + 1 # 统计不同厂商的设备数量 vendor_stats = {} for node in topology.nodes.values(): vendor = node.vendor vendor_stats[vendor] = vendor_stats.get(vendor, 0) + 1 return { "success": True, "total_nodes": len(topology.nodes), "total_links": len(topology.links), "protocol_distribution": protocol_stats, "vendor_distribution": vendor_stats, "last_discovery": topology.last_discovery.isoformat() if topology.last_discovery else None, "discovery_scope": topology.discovery_scope } except Exception as e: logger.error(f"获取拓扑统计失败: {e}") return {"success": False, "message": f"获取拓扑统计失败: {str(e)}"} # 工具分类：网络扫描 @mcp.tool() async def scan_network_range( network: str, timeout: float = 3.0, max_concurrent: int = 50, ping_enabled: bool = True, port_scan_enabled: bool = True, snmp_enabled: bool = True, auto_create_devices: bool = False ) -> Dict[str, Any]: """ 扫描指定网络范围内的活跃设备 Args: network: 网络范围，如 "192.168.1.0/24" 或 "10.0.0.0/16" timeout: 扫描超时时间（秒） max_concurrent: 最大并发扫描数 ping_enabled: 是否启用ping扫描 port_scan_enabled: 是否启用端口扫描 snmp_enabled: 是否启用SNMP扫描 auto_create_devices: 是否自动创建发现的设备 Returns: 扫描结果 """ try: from network_scanner import ScanConfiguration # 创建扫描配置 config = ScanConfiguration( timeout=timeout, max_concurrent=max_concurrent, ping_enabled=ping_enabled, port_scan_enabled=port_scan_enabled, snmp_enabled=snmp_enabled ) logger.info(f"开始网络扫描: {network}") # 执行网络扫描 scan_results = await network_scanner.scan_network_range(network, config) # 如果启用自动创建设备，则创建发现的设备 discovered_devices = [] if auto_create_devices: discovered_devices = await network_scanner.discover_devices_from_scan(scan_results, auto_create=True) return { "success": True, "message": f"网络扫描完成，发现 {len(scan_results)} 个活跃主机", "network_range": network, "alive_hosts": len(scan_results), "scan_results": [result.to_dict() for result in scan_results], "auto_created_devices": len(discovered_devices), "discovered_devices": [device.to_dict() for device in discovered_devices] } except Exception as e: logger.error(f"网络扫描失败: {e}") return {"success": False, "message": f"网络扫描失败: {str(e)}"} @mcp.tool() async def get_scan_results( ip_address: Optional[str] = None, alive_only: bool = True ) -> Dict[str, Any]: """ 获取网络扫描结果 Args: ip_address: 特定IP地址，如果不提供则返回所有结果 alive_only: 是否只返回活跃主机 Returns: 扫描结果 """ try: if ip_address: # 获取特定IP的扫描结果 result = network_scanner.scan_results.get(ip_address) if result: return { "success": True, "scan_result": result.to_dict() } else: return {"success": False, "message": f"未找到IP {ip_address} 的扫描结果"} else: # 获取所有扫描结果 results = list(network_scanner.scan_results.values()) if alive_only: results = [r for r in results if r.is_alive] return { "success": True, "total_results": len(results), "scan_results": [result.to_dict() for result in results] } except Exception as e: logger.error(f"获取扫描结果失败: {e}") return {"success": False, "message": f"获取扫描结果失败: {str(e)}"} @mcp.tool() async def get_scan_statistics() -> Dict[str, Any]: """ 获取网络扫描统计信息 Returns: 扫描统计数据 """ try: stats = network_scanner.get_scan_statistics() return { "success": True, "statistics": stats } except Exception as e: logger.error(f"获取扫描统计失败: {e}") return {"success": False, "message": f"获取扫描统计失败: {str(e)}"} @mcp.tool() async def discover_devices_from_scan_results( min_response_time: Optional[float] = None, required_ports: str = "", vendor_filter: Optional[str] = None, device_type_filter: Optional[str] = None ) -> Dict[str, Any]: """ 从扫描结果中发现并创建设备 Args: min_response_time: 最小响应时间过滤（毫秒） required_ports: 必需的开放端口，使用逗号分隔，如"22,23,161" vendor_filter: 厂商过滤，如"cisco"、"huawei" device_type_filter: 设备类型过滤，如"switch"、"router" Returns: 发现和创建的设备信息 """ try: # 获取所有活跃的扫描结果 all_results = [r for r in network_scanner.scan_results.values() if r.is_alive] # 应用过滤条件 filtered_results = [] for result in all_results: # 响应时间过滤 if min_response_time and result.response_time and result.response_time > min_response_time: continue # 端口过滤 if required_ports: required_port_list = [int(p.strip()) for p in required_ports.split(',') if p.strip().isdigit()] if not all(port in result.open_ports for port in required_port_list): continue # 厂商过滤 if vendor_filter and result.vendor: if vendor_filter.lower() not in result.vendor.lower(): continue # 设备类型过滤 if device_type_filter and result.device_type: if device_type_filter.lower() not in result.device_type.lower(): continue filtered_results.append(result) # 从过滤后的结果创建设备 discovered_devices = await network_scanner.discover_devices_from_scan(filtered_results, auto_create=True) return { "success": True, "message": f"从 {len(filtered_results)} 个扫描结果中创建了 {len(discovered_devices)} 个设备", "filtered_results_count": len(filtered_results), "created_devices_count": len(discovered_devices), "created_devices": [device.to_dict() for device in discovered_devices] } except Exception as e: logger.error(f"从扫描结果创建设备失败: {e}") return {"success": False, "message": f"从扫描结果创建设备失败: {str(e)}"} @mcp.tool() async def clear_scan_results() -> Dict[str, Any]: """ 清空网络扫描结果 Returns: 操作结果 """ try: network_scanner.clear_scan_results() return {"success": True, "message": "扫描结果已清空"} except Exception as e: logger.error(f"清空扫描结果失败: {e}") return {"success": False, "message": f"清空扫描结果失败: {str(e)}"} @mcp.tool() async def scan_single_host( ip_address: str, timeout: float = 3.0, port_scan_enabled: bool = True, snmp_enabled: bool = True ) -> Dict[str, Any]: """ 扫描单个主机 Args: ip_address: 目标IP地址 timeout: 扫描超时时间（秒） port_scan_enabled: 是否启用端口扫描 snmp_enabled: 是否启用SNMP扫描 Returns: 单个主机扫描结果 """ try: from network_scanner import ScanConfiguration # 创建扫描配置 config = ScanConfiguration( timeout=timeout, port_scan_enabled=port_scan_enabled, snmp_enabled=snmp_enabled ) logger.info(f"开始扫描单个主机: {ip_address}") # 执行单个主机扫描 scan_result = await network_scanner.scan_single_host(ip_address, config) # 保存到扫描结果 network_scanner.scan_results[ip_address] = scan_result network_scanner.save_scan_results() return { "success": True, "message": f"主机扫描完成: {ip_address}", "is_alive": scan_result.is_alive, "scan_result": scan_result.to_dict() } except Exception as e: logger.error(f"扫描单个主机失败: {e}") return {"success": False, "message": f"扫描单个主机失败: {str(e)}"} # =================== 资源工具 =================== @mcp.tool() async def list_resources() -> List[Dict[str, Any]]: """ 列出可用的MCP资源 Returns: 资源列表 """ return resource_manager.list_available_resources() @mcp.tool() async def get_resource(uri: str, use_cache: bool = True) -> Dict[str, Any]: """ 获取指定URI的资源 Args: uri: 资源URI use_cache: 是否使用缓存 Returns: 资源内容 """ return await resource_manager.get_resource(uri, use_cache) @mcp.tool() async def clear_resource_cache(uri: Optional[str] = None) -> Dict[str, Any]: """ 清除资源缓存 Args: uri: 要清除的特定资源URI，如果为None则清除所有缓存 Returns: 操作结果 """ success = resource_manager.clear_cache(uri) if success: if uri: return {"success": True, "message": f"已清除资源缓存: {uri}"} else: return {"success": True, "message": "已清除所有资源缓存"} else: return {"success": False, "message": "清除缓存失败"} # =================== 模板工具 =================== @mcp.tool() async def list_templates() -> List[Dict[str, Any]]: """ 列出可用的提示模板 Returns: 模板列表 """ return template_manager.list_templates() @mcp.tool() async def render_template( template_name: str, context: str = "{}", resource_uris: str = "" ) -> Dict[str, Any]: """ 渲染提示模板 Args: template_name: 模板名称 context: JSON格式的渲染上下文，如 {"name": "value"} resource_uris: 资源URI映射，格式为 "key1=uri1,key2=uri2" Returns: 渲染结果 """ try: # 解析context context_dict = {} if context.strip(): try: context_dict = json.loads(context) if not isinstance(context_dict, dict): context_dict = {} except json.JSONDecodeError: logger.warning(f"无法解析context JSON: {context}") return {"success": False, "message": "无效的JSON上下文格式"} # 解析resource_uris resource_dict = {} if resource_uris.strip(): uri_pairs = resource_uris.split(',') for pair in uri_pairs: if '=' in pair: key, uri = pair.split('=', 1) resource_dict[key.strip()] = uri.strip() if resource_dict: # 渲染带有资源的模板 result = await render_template_with_resources( template_name=template_name, context=context_dict, resource_uris=resource_dict, resource_manager=resource_manager ) else: # 渲染简单模板 result = template_manager.render_template(template_name, context_dict) if result: if isinstance(result, list): # 处理消息列表 messages = [ { "role": msg.role, "content": msg.content } for msg in result ] return {"success": True, "messages": messages} else: # 处理字符串结果 return {"success": True, "rendered_template": result} else: return {"success": False, "message": f"渲染模板失败: {template_name}"} except Exception as e: logger.error(f"渲染模板异常: {str(e)}") return {"success": False, "message": f"渲染模板异常: {str(e)}"} # =================== 资源注册 =================== # 将资源处理函数注册到资源管理器 @resource_manager.register_resource("greeting/{name}", ResourceType.SYSTEM) async def get_greeting(name: str) -> str: """获取个性化问候语""" return f"你好，{name}！欢迎使用NetBrain MCP！" @resource_manager.register_resource("device/{device_id}", ResourceType.DEVICE) async def get_device_resource(device_id: str) -> Dict[str, Any]: """获取设备资源""" device = device_manager.get_device(device_id) if device: return device.to_dict() return {"error": f"设备不存在: {device_id}"} @resource_manager.register_resource("device/{device_id}/config", ResourceType.CONFIG) async def get_device_config(device_id: str) -> Dict[str, Any]: """获取设备配置资源""" device = device_manager.get_device(device_id) if not device: return {"error": f"设备不存在: {device_id}"} # 这里我们会尝试获取设备的运行配置 # 首先获取默认凭据 credentials = device_manager.list_credentials() if not credentials: return {"error": "没有可用的凭据"} credential_id = device.credential_id or credentials[0].id credential = device_manager.get_credential(credential_id) # 尝试获取设备配置 if device.vendor == DeviceVendor.CISCO: command = "show running-config" elif device.vendor == DeviceVendor.HUAWEI: command = "display current-configuration" elif device.vendor == DeviceVendor.H3C: command = "display current-configuration" elif device.vendor == DeviceVendor.JUNIPER: command = "show configuration" else: command = "show running-config" # 默认命令 try: # 连接设备 success, error = await connection_manager.connect_device(device, credential) if not success: return {"error": f"连接设备失败: {error}"} # 发送命令 result, error = await connection_manager.send_command( device_id=device_id, credential_id=credential_id, command=command, timeout=60 # 配置命令可能需要更长时间 ) # 断开连接 await connection_manager.disconnect_device(device_id, credential_id) if result and result.success: return { "device_id": device_id, "device_name": device.name, "config_type": "running", "config_content": result.output, "timestamp": result.execution_time.isoformat() } else: return {"error": f"获取配置失败: {error or '未知错误'}"} except Exception as e: return {"error": f"获取配置时发生错误: {str(e)}"} @resource_manager.register_resource("device/{device_id}/interfaces", ResourceType.DEVICE) async def get_device_interfaces(device_id: str) -> Dict[str, Any]: """获取设备接口资源""" device = device_manager.get_device(device_id) if not device: return {"error": f"设备不存在: {device_id}"} # 获取默认凭据 credentials = device_manager.list_credentials() if not credentials: return {"error": "没有可用的凭据"} credential_id = device.credential_id or credentials[0].id credential = device_manager.get_credential(credential_id) # 根据设备类型确定接口命令 if device.vendor == DeviceVendor.CISCO: command = "show interfaces" elif device.vendor == DeviceVendor.HUAWEI: command = "display interface" elif device.vendor == DeviceVendor.H3C: command = "display interface" elif device.vendor == DeviceVendor.JUNIPER: command = "show interfaces detail" else: command = "show interfaces" # 默认命令 try: # 连接设备 success, error = await connection_manager.connect_device(device, credential) if not success: return {"error": f"连接设备失败: {error}"} # 发送命令 result, error = await connection_manager.send_command( device_id=device_id, credential_id=credential_id, command=command, timeout=30 ) # 断开连接 await connection_manager.disconnect_device(device_id, credential_id) if result and result.success: return { "device_id": device_id, "device_name": device.name, "interfaces_output": result.output, "timestamp": result.execution_time.isoformat() } else: return {"error": f"获取接口信息失败: {error or '未知错误'}"} except Exception as e: return {"error": f"获取接口信息时发生错误: {str(e)}"} @resource_manager.register_resource("device/{device_id}/routes", ResourceType.DEVICE) async def get_device_routes(device_id: str) -> Dict[str, Any]: """获取设备路由表资源""" device = device_manager.get_device(device_id) if not device: return {"error": f"设备不存在: {device_id}"} # 获取默认凭据 credentials = device_manager.list_credentials() if not credentials: return {"error": "没有可用的凭据"} credential_id = device.credential_id or credentials[0].id credential = device_manager.get_credential(credential_id) # 根据设备类型确定路由命令 if device.vendor == DeviceVendor.CISCO: command = "show ip route" elif device.vendor == DeviceVendor.HUAWEI: command = "display ip routing-table" elif device.vendor == DeviceVendor.H3C: command = "display ip routing-table" elif device.vendor == DeviceVendor.JUNIPER: command = "show route" else: command = "show ip route" # 默认命令 try: # 连接设备 success, error = await connection_manager.connect_device(device, credential) if not success: return {"error": f"连接设备失败: {error}"} # 发送命令 result, error = await connection_manager.send_command( device_id=device_id, credential_id=credential_id, command=command, timeout=30 ) # 断开连接 await connection_manager.disconnect_device(device_id, credential_id) if result and result.success: return { "device_id": device_id, "device_name": device.name, "routes_output": result.output, "timestamp": result.execution_time.isoformat() } else: return {"error": f"获取路由表失败: {error or '未知错误'}"} except Exception as e: return {"error": f"获取路由表时发生错误: {str(e)}"} @resource_manager.register_resource("credentials", ResourceType.CREDENTIAL) async def list_credentials_resource() -> Dict[str, Any]: """ 列出所有设备凭据 Returns: 凭据列表 """ credentials = device_manager.list_credentials() return { "credentials": [ { "id": cred.id, "name": cred.name, "username": cred.username, "password": cred.password, "protocol": cred.protocol.value, "port": cred.port, "ssh_key_file": cred.ssh_key_file } for cred in credentials ] } @resource_manager.register_resource("system/status", ResourceType.SYSTEM) async def get_system_status() -> Dict[str, Any]: """获取系统状态资源""" devices_count = len(device_manager.list_devices()) credentials_count = len(device_manager.list_credentials()) active_connections = len(connection_manager.get_active_connections()) return { "status": "running", "devices_count": devices_count, "credentials_count": credentials_count, "active_connections": active_connections, "version": "0.1.0" } @resource_manager.register_resource("topology", ResourceType.TOPOLOGY) async def get_topology_resource() -> Dict[str, Any]: """获取网络拓扑资源""" try: topology = topology_discovery.get_topology() return topology.to_dict() except Exception as e: logger.error(f"获取拓扑资源失败: {e}") return {"error": f"获取拓扑资源失败: {str(e)}"} @resource_manager.register_resource("topology/statistics", ResourceType.TOPOLOGY) async def get_topology_statistics_resource() -> Dict[str, Any]: """获取拓扑统计资源""" try: topology = topology_discovery.get_topology() # 统计不同协议的链路数量 protocol_stats = {} for link in topology.links: protocol = link.protocol.value protocol_stats[protocol] = protocol_stats.get(protocol, 0) + 1 # 统计不同厂商的设备数量 vendor_stats = {} for node in topology.nodes.values(): vendor = node.vendor vendor_stats[vendor] = vendor_stats.get(vendor, 0) + 1 return { "total_nodes": len(topology.nodes), "total_links": len(topology.links), "protocol_distribution": protocol_stats, "vendor_distribution": vendor_stats, "last_discovery": topology.last_discovery.isoformat() if topology.last_discovery else None, "discovery_scope": topology.discovery_scope } except Exception as e: logger.error(f"获取拓扑统计资源失败: {e}") return {"error": f"获取拓扑统计资源失败: {str(e)}"} @resource_manager.register_resource("device/{device_id}/neighbors", ResourceType.TOPOLOGY) async def get_device_neighbors_resource(device_id: str) -> Dict[str, Any]: """获取设备邻居资源""" try: topology = topology_discovery.get_topology() neighbors = topology.get_device_neighbors(device_id) # 获取邻居设备的详细信息 neighbor_details = [] for neighbor_id in neighbors: device = device_manager.get_device(neighbor_id) if device: neighbor_details.append(device.to_dict()) return { "device_id": device_id, "neighbors": neighbor_details, "neighbor_count": len(neighbor_details) } except Exception as e: logger.error(f"获取设备邻居资源失败: {e}") return {"error": f"获取设备邻居资源失败: {str(e)}"} @resource_manager.register_resource("scan/results", ResourceType.SCAN) async def get_scan_results_resource() -> Dict[str, Any]: """获取网络扫描结果资源""" try: results = [r for r in network_scanner.scan_results.values() if r.is_alive] return { "total_results": len(results), "scan_results": [result.to_dict() for result in results] } except Exception as e: logger.error(f"获取扫描结果资源失败: {e}") return {"error": f"获取扫描结果资源失败: {str(e)}"} @resource_manager.register_resource("scan/statistics", ResourceType.SCAN) async def get_scan_statistics_resource() -> Dict[str, Any]: """获取网络扫描统计资源""" try: stats = network_scanner.get_scan_statistics() return stats except Exception as e: logger.error(f"获取扫描统计资源失败: {e}") return {"error": f"获取扫描统计资源失败: {str(e)}"} @resource_manager.register_resource("scan/result/{ip_address}", ResourceType.SCAN) async def get_scan_result_resource(ip_address: str) -> Dict[str, Any]: """获取特定IP的扫描结果资源""" try: result = network_scanner.scan_results.get(ip_address) if result: return result.to_dict() else: return {"error": f"未找到IP {ip_address} 的扫描结果"} except Exception as e: logger.error(f"获取扫描结果资源失败: {e}") return {"error": f"获取扫描结果资源失败: {str(e)}"} # 向MCP服务器注册资源获取方法 @mcp.resource("{uri}") async def mcp_resource_handler(uri: str) -> Any: """MCP资源处理器，处理所有资源请求""" logger.info(f"处理资源请求: {uri}") return await resource_manager.get_resource(uri) # 向MCP服务器注册提示模板 @mcp.prompt("{name}") async def mcp_prompt_handler(name: str) -> str: """MCP提示模板处理器，处理所有模板请求""" logger.info(f"处理模板请求: {name}") template = template_manager.get_template(name) if template: # 返回模板描述或模板名称 return template.description or template.name return f"错误: 未找到模板 '{name}'" # 测试工具 @mcp.tool() async def test_scrapli_connection( host: str, username: str, password: str = None, platform: str = "cisco_iosxe", port: int = 22, protocol: str = "ssh", ssh_key_file: Optional[str] = None, connect_timeout: int = 15 ) -> Dict[str, Any]: """ 测试Scrapli连接（直接使用Scrapli） Args: host: 设备主机名或IP地址 username: 用户名 password: 密码（与ssh_key_file至少提供一个） platform: 平台类型，如cisco_iosxe, huawei_vrp等 port: 端口号（默认22，Telnet使用23） protocol: 连接协议，ssh或telnet ssh_key_file: SSH密钥文件路径，可选 connect_timeout: 连接超时时间（秒） Returns: 连接测试结果 """ # 从device_connector中导入SCRAPLI_IMPORT_SUCCESS检查scrapli是否可用 if not SCRAPLI_IMPORT_SUCCESS: return { "success": False, "message": "Scrapli库导入失败，请确保已安装: pip install scrapli scrapli-community" } try: from scrapli import Scrapli # 修正平台名称处理 if platform.lower() == "huawei": platform = "huawei_vrp" # 确定连接协议 is_telnet = protocol.lower() == "telnet" if is_telnet and port == 22: # 如果是Telnet但端口仍为默认SSH端口，修正为默认Telnet端口 port = 23 logger.info(f"连接参数: host={host}, username={username}, platform={platform}, port={port}, protocol={protocol}") # 打印密码信息（安全地） if password: logger.info(f"密码类型: {type(password)}, 长度: {len(str(password))}") password = str(password) # 确保密码是字符串类型 # 准备连接参数 device_params = { "host": host, "auth_username": username, "auth_strict_key": False, "platform": platform, "port": port, "timeout_socket": connect_timeout, "timeout_transport": connect_timeout * 2, } # 根据协议设置传输方式 if is_telnet: device_params["transport"] = "telnet" logger.info("使用Telnet连接") else: # 指定Windows兼容的传输方式 device_params["transport"] = "paramiko" # 在Windows上使用paramiko而不是system logger.info("使用SSH连接 (paramiko)") # 设置认证方式，优先使用SSH密钥 if not is_telnet and ssh_key_file and os.path.exists(ssh_key_file): device_params["auth_private_key"] = ssh_key_file logger.info(f"使用SSH密钥认证: {ssh_key_file}") elif password: device_params["auth_password"] = password logger.info("使用密码认证") else: return {"success": False, "message": "必须提供密码或SSH密钥文件"} # 详细记录连接参数（移除敏感信息） safe_params = device_params.copy() if "auth_password" in safe_params: safe_params["auth_password"] = "******" if "auth_private_key" in safe_params: safe_params["auth_private_key"] = f"[使用密钥文件: {ssh_key_file}]" logger.info(f"Scrapli测试连接参数: {safe_params}") # 建立连接 conn = Scrapli(**device_params) conn.open() # 获取设备提示符 prompt = conn.get_prompt() # 测试发送命令 try: # 根据平台类型选择命令 if "huawei" in platform: command = "display version" elif "cisco" in platform: command = "show version" else: command = "show version" resp = conn.send_command(command) version_info = resp.result[:200] + "..." if len(resp.result) > 200 else resp.result except Exception as e: logger.warning(f"获取版本信息失败: {str(e)}") version_info = "无法获取版本信息" # 关闭连接 conn.close() return { "success": True, "message": f"连接成功！设备提示符: {prompt}", "prompt": prompt, "version_info": version_info } except ImportError as e: return {"success": False, "message": f"Scrapli库未安装: {str(e)}"} except Exception as e: error_type = type(e).__name__ error_msg = str(e) logger.error(f"Scrapli连接测试失败: {error_type} - {error_msg}") return { "success": False, "message": f"连接失败: {error_type} - {error_msg}", "error_type": error_type, "error_details": error_msg } @mcp.tool() async def test_telnet_connection( host: str, username: str, password: str, platform: str = "cisco_iosxe", port: int = 23, connect_timeout: int = 15 ) -> Dict[str, Any]: """ 测试Telnet连接（快速测试） Args: host: 设备主机名或IP地址 username: 用户名 password: 密码 platform: 平台类型，如cisco_iosxe, huawei_vrp等 port: 端口号（默认23） connect_timeout: 连接超时时间（秒） Returns: 连接测试结果 """ # 直接调用test_scrapli_connection，但固定protocol为telnet return await test_scrapli_connection( host=host, username=username, password=password, platform=platform, port=port, protocol="telnet", connect_timeout=connect_timeout ) @mcp.tool() async def send_telnet_command( device_id: str, credential_id: str, command: str, timeout: int = 60 ) -> Dict[str, Any]: """ 向Telnet连接的网络设备发送命令（使用更长的超时时间） Args: device_id: 设备ID credential_id: 凭据ID command: 要执行的命令 timeout: 命令超时时间（秒，默认比普通命令更长） Returns: 命令执行结果 """ # 获取设备和凭据 device = device_manager.get_device(device_id) credential = device_manager.get_credential(credential_id) if not device: return {"success": False, "message": f"设备不存在: {device_id}"} if not credential: return {"success": False, "message": f"凭据不存在: {credential_id}"} # 验证是否为Telnet凭据 if credential.protocol != ConnectionProtocol.TELNET: return {"success": False, "message": f"此函数仅支持Telnet连接，当前凭据协议为: {credential.protocol.value}"} # 记录命令执行信息 logger.info(f"通过Telnet向设备 {device.name} ({device.ip_address}) 发送命令: {command}") logger.info(f"超时设置: {timeout}秒") # 发送命令 result, error = await connection_manager.send_command( device_id=device_id, credential_id=credential_id, command=command, timeout=timeout ) if result: # 记录执行结果 if result.success: logger.info(f"命令执行成功，输出长度: {len(result.output)}") else: logger.warning(f"命令执行失败: {result.error_message}") return { "success": result.success, "command": result.command, "output": result.output, "error_message": result.error_message, "execution_time": result.execution_time.isoformat() } else: logger.error(f"执行命令失败: {error}") return { "success": False, "command": command, "output": "", "error_message": error, "execution_time": None } # 运行服务器 if __name__ == "__main__": # 检测启动方式 import os # 获取启动命令行参数 import sys command_args = " ".join(sys.argv) # 检查环境变量或命令行，判断是否通过MCP CLI运行 is_mcp_cli = os.environ.get("MCP_CLI_RUN", "0") == "1" or "mcp dev" in command_args or "mcp run" in command_args # 设置环境变量，使其他模块也能知道是通过MCP CLI运行 if is_mcp_cli: os.environ["MCP_CLI_RUN"] = "1" # 始终启动Web服务器（后台运行） import threading import uvicorn def start_web_server(): """启动Web服务器""" from web.app import app uvicorn.run(app, host="0.0.0.0", port=8088) # 在另一个线程中启动Web服务器 web_thread = threading.Thread(target=start_web_server, daemon=True) web_thread.start() print("NetBrain MCP Web界面已启动，访问 http://localhost:8088") if is_mcp_cli: print("MCP CLI模式已检测到") # 在8000端口运行，与MCP Inspector预期一致 mcp.run()