server.py•60.1 kB
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()