core_modules.md•21.2 kB
# NetBrain MCP 核心模块设计
本文档详细描述NetBrain MCP系统的核心模块设计，包括模块功能、接口以及之间的交互关系。
## 1. MCP服务器模块
### 1.1 模块概述
MCP服务器模块是整个系统的入口点，负责处理来自大型语言模型的请求，并将其路由到相应的工具或资源处理器。该模块基于FastMCP框架实现，封装了MCP协议的实现细节，提供了统一的API接口。
**实现文件**: `server.py`
### 1.2 核心类与接口
#### 1.2.1 FastMCP服务器
```python
# 创建MCP服务器实例
mcp = FastMCP("NetBrain MCP")
# 工具注册装饰器
@mcp.tool()
async def tool_function(param1: str, param2: int) -> Dict[str, Any]:
"""工具功能实现"""
pass
# 资源注册装饰器
@mcp.resource("{uri}")
async def mcp_resource_handler(uri: str) -> Any:
"""MCP资源处理器，处理所有资源请求"""
return await resource_manager.get_resource(uri)
# 提示模板注册装饰器
@mcp.prompt("{name}")
async def mcp_prompt_handler(name: str) -> str:
"""MCP提示模板处理器"""
return await template_manager.render_template(name)
```
### 1.3 功能特性
- **工具注册与执行**：目前已实现32个MCP工具，涵盖设备管理、连接管理、拓扑发现等功能
- **资源注册与提供**：实现13个MCP资源，包括设备、配置、拓扑等资源类型
- **提示模板管理**：集成模板系统，支持动态模板渲染
- **会话管理**：维护客户端会话状态和连接信息
- **错误处理**：统一的错误处理机制和日志记录
### 1.4 与其他模块的交互
- 与工具管理模块交互，通过装饰器方式注册工具
- 与资源模块交互，通过resource_manager获取和管理资源
- 与模板系统交互，通过template_manager渲染提示模板
- 与设备管理和连接模块交互，提供网络设备操作功能
## 2. 工具管理模块
### 2.1 模块概述
工具管理模块负责工具的注册、分类和执行管理，虽然MCP工具直接通过FastMCP注册，但该模块提供了工具的分类和管理功能。
**实现文件**: `tool_manager.py`
### 2.2 核心类与接口
#### 2.2.1 ToolManager类
```python
class ToolManager:
"""工具管理器，负责工具的注册和调用"""
def __init__(self):
self.tools: Dict[str, ToolInfo] = {}
self.categories: Dict[ToolCategory, List[str]] = {cat: [] for cat in ToolCategory}
def register_tool(self,
name: Optional[str] = None,
description: Optional[str] = None,
category: ToolCategory = ToolCategory.GENERAL) -> Callable:
"""工具注册装饰器"""
def list_tools(self, category: Optional[ToolCategory] = None) -> List[Dict[str, Any]]:
"""列出已注册的工具"""
def get_tool(self, name: str) -> Optional[ToolInfo]:
"""获取工具信息"""
def execute_tool(self, name: str, *args, **kwargs) -> Any:
"""执行工具"""
```
#### 2.2.2 ToolCategory枚举
```python
class ToolCategory(Enum):
"""工具分类枚举"""
GENERAL = "general" # 通用工具
NETWORK_DEVICE = "network_device" # 网络设备相关工具
CONFIGURATION = "configuration" # 配置管理工具
TOPOLOGY = "topology" # 拓扑相关工具
DIAGNOSTIC = "diagnostic" # 诊断工具
SECURITY = "security" # 安全相关工具
```
#### 2.2.3 ToolInfo数据类
```python
@dataclass
class ToolInfo:
"""工具信息数据类"""
name: str # 工具名称
func: Callable # 工具函数
description: str # 工具描述
category: ToolCategory # 工具分类
parameters: Dict[str, Dict[str, Any]] # 参数信息
return_type: str # 返回类型
```
### 2.3 实际工具实现
系统中实际实现的32个MCP工具分类如下：
#### 设备管理工具 (6个)
- `list_devices` - 列出网络设备
- `add_device` - 添加新的网络设备
- `get_device` - 获取设备详细信息
- `update_device` - 更新设备信息
- `delete_device` - 删除设备
#### 凭据管理工具 (2个)
- `add_credential` - 添加设备凭据
- `list_credentials` - 列出所有设备凭据
#### 设备连接工具 (4个)
- `connect_device` - 连接到网络设备
- `disconnect_device` - 断开与网络设备的连接
- `send_command` - 发送单个命令
- `send_commands` - 发送多个命令
- `get_active_connections` - 获取活动连接
#### 拓扑发现工具 (6个)
- `discover_topology` - 发现网络拓扑
- `get_topology` - 获取网络拓扑
- `clear_topology` - 清除拓扑缓存
- `get_device_neighbors` - 获取设备邻居
- `discover_device_neighbors` - 发现设备邻居
- `get_topology_statistics` - 获取拓扑统计
#### 网络扫描工具 (5个)
- `scan_network_range` - 扫描网络范围
- `get_scan_results` - 获取扫描结果
- `get_scan_statistics` - 获取扫描统计
- `discover_devices_from_scan_results` - 从扫描结果发现设备
- `clear_scan_results` - 清除扫描结果
- `scan_single_host` - 扫描单个主机
#### 资源管理工具 (3个)
- `list_resources` - 列出可用资源
- `get_resource` - 获取资源内容
- `clear_resource_cache` - 清除资源缓存
#### 模板管理工具 (2个)
- `list_templates` - 列出可用模板
- `render_template` - 渲染模板
#### 测试工具 (4个)
- `test_scrapli_connection` - 测试Scrapli连接
- `test_telnet_connection` - 测试Telnet连接
- `send_telnet_command` - 发送Telnet命令
## 3. 资源模块
### 3.1 模块概述
资源模块负责管理和提供各种网络资源，资源是客户端可以访问的数据对象，通过URI标识。系统实现了13个MCP资源。
**实现文件**: `mcp_resources.py`
### 3.2 核心类与接口
#### 3.2.1 ResourceManager类
```python
class ResourceManager:
"""MCP资源管理器，负责资源注册和提供"""
def __init__(self):
self.resources = {}
self.resource_patterns = {}
self.resource_cache = {}
self.cache_expiration = {}
self.default_cache_ttl = 300 # 默认缓存5分钟
self.cache_dir = os.path.join(os.getcwd(), "resource_cache")
def register_resource(self, uri_pattern: str, resource_type: str):
"""注册资源装饰器"""
async def get_resource(self, uri: str, use_cache: bool = True, cache_ttl: Optional[int] = None) -> Dict[str, Any]:
"""获取资源"""
def clear_cache(self, uri: Optional[str] = None) -> bool:
"""清除资源缓存"""
def list_available_resources(self) -> List[Dict[str, Any]]:
"""列出可用的资源"""
```
#### 3.2.2 ResourceType类
```python
class ResourceType:
"""资源类型定义"""
DEVICE = "device" # 设备资源
CREDENTIAL = "credential" # 凭据资源
TOPOLOGY = "topology" # 拓扑资源
SYSTEM = "system" # 系统资源
CONFIG = "config" # 配置资源
LOG = "log" # 日志资源
REPORT = "report" # 报告资源
SCAN = "scan" # 扫描资源
```
### 3.3 实际资源实现
系统中实际实现的13个MCP资源：
1. `greeting/{name}` - 问候信息（示例资源）
2. `device/{device_id}` - 设备基本信息
3. `device/{device_id}/config` - 设备配置信息
4. `device/{device_id}/interfaces` - 设备接口信息
5. `device/{device_id}/routes` - 设备路由信息
6. `credentials` - 凭据列表
7. `system/status` - 系统状态
8. `topology` - 网络拓扑
9. `topology/statistics` - 拓扑统计
10. `device/{device_id}/neighbors` - 设备邻居
11. `scan/results` - 扫描结果
12. `scan/statistics` - 扫描统计
13. `scan/result/{ip_address}` - 单个IP扫描结果
## 4. 设备管理模块
### 4.1 模块概述
设备管理模块负责网络设备的生命周期管理，包括设备的添加、查询、更新和删除，以及设备凭据的管理。
**实现文件**: `network_devices.py`
### 4.2 核心类与接口
#### 4.2.1 DeviceManager类
```python
class DeviceManager:
"""设备管理器，负责设备的创建、查询和管理"""
def __init__(self):
self.devices: Dict[str, NetworkDevice] = {}
self.credentials: Dict[str, DeviceCredential] = {}
self.data_dir = "data"
self.devices_file = os.path.join(self.data_dir, "devices.json")
self.credentials_file = os.path.join(self.data_dir, "credentials.json")
def add_device(self, device: NetworkDevice) -> str:
"""添加设备"""
def get_device(self, device_id: str) -> Optional[NetworkDevice]:
"""获取设备"""
def update_device(self, device_id: str, **kwargs) -> Optional[NetworkDevice]:
"""更新设备信息"""
def delete_device(self, device_id: str) -> bool:
"""删除设备"""
def list_devices(self,
vendor: Optional[DeviceVendor] = None,
device_type: Optional[DeviceType] = None,
status: Optional[DeviceStatus] = None,
tag: Optional[str] = None) -> List[NetworkDevice]:
"""列出设备，支持过滤"""
def add_credential(self, credential: DeviceCredential) -> str:
"""添加凭据"""
def get_credential(self, credential_id: str) -> Optional[DeviceCredential]:
"""获取凭据"""
def delete_credential(self, credential_id: str) -> bool:
"""删除凭据"""
def list_credentials(self) -> List[DeviceCredential]:
"""列出所有凭据"""
```
#### 4.2.2 网络设备相关枚举
```python
class DeviceType(Enum):
"""网络设备类型枚举"""
ROUTER = "router"
SWITCH = "switch"
FIREWALL = "firewall"
LOAD_BALANCER = "load_balancer"
WIRELESS_CONTROLLER = "wireless_controller"
ACCESS_POINT = "access_point"
OTHER = "other"
class DeviceVendor(Enum):
"""网络设备厂商枚举"""
CISCO = "cisco"
HUAWEI = "huawei"
H3C = "h3c"
JUNIPER = "juniper"
ARISTA = "arista"
FORTINET = "fortinet"
CHECKPOINT = "checkpoint"
OTHER = "other"
class DeviceStatus(Enum):
"""设备状态枚举"""
ONLINE = "online"
OFFLINE = "offline"
UNREACHABLE = "unreachable"
MAINTENANCE = "maintenance"
UNKNOWN = "unknown"
class ConnectionProtocol(Enum):
"""连接协议枚举"""
SSH = "ssh"
TELNET = "telnet"
SNMP = "snmp"
HTTP = "http"
HTTPS = "https"
NETCONF = "netconf"
```
#### 4.2.3 设备数据模型
```python
@dataclass
class NetworkDevice:
"""网络设备模型"""
id: str = field(default_factory=lambda: str(uuid.uuid4()))
name: str = ""
ip_address: str = ""
device_type: DeviceType = DeviceType.OTHER
vendor: DeviceVendor = DeviceVendor.OTHER
platform: str = "" # Scrapli平台类型
model: str = ""
os_version: str = ""
status: DeviceStatus = DeviceStatus.UNKNOWN
location: str = ""
credential_id: Optional[str] = None
description: str = ""
tags: List[str] = field(default_factory=list)
last_seen: Optional[datetime] = None
created_at: datetime = field(default_factory=datetime.now)
updated_at: datetime = field(default_factory=datetime.now)
custom_attributes: Dict[str, Any] = field(default_factory=dict)
```
#### 4.2.4 凭据数据模型
```python
class DeviceCredential:
"""设备认证凭据"""
def __init__(self,
id: str = None,
name: str = "",
username: str = "",
password: str = "",
protocol: ConnectionProtocol = ConnectionProtocol.SSH,
port: int = None,
enable_password: str = None,
ssh_key_file: str = None):
# 实现细节...
```
## 5. 设备连接模块
### 5.1 模块概述
设备连接模块负责与网络设备建立连接并执行命令。该模块基于Scrapli库实现，支持SSH、Telnet等多种连接协议，并提供统一的命令执行接口。
**实现文件**: `device_connector.py`
### 5.2 核心类与接口
#### 5.2.1 ConnectionManager类
```python
class ConnectionManager:
"""连接管理器，负责管理所有设备连接"""
def __init__(self):
self.active_connections: Dict[str, DeviceConnector] = {}
async def connect_device(self,
device: NetworkDevice,
credential: DeviceCredential) -> Tuple[bool, Optional[str]]:
"""连接设备"""
async def disconnect_device(self, device_id: str, credential_id: str) -> Tuple[bool, Optional[str]]:
"""断开设备连接"""
async def send_command(self,
device_id: str,
credential_id: str,
command: str,
timeout: int = 30) -> Tuple[Optional[CommandResult], Optional[str]]:
"""发送命令到设备"""
async def send_commands(self,
device_id: str,
credential_id: str,
commands: List[str],
timeout: int = 30) -> Tuple[Optional[List[CommandResult]], Optional[str]]:
"""发送多个命令到设备"""
def get_active_connections(self) -> List[Dict[str, Any]]:
"""获取活动连接列表"""
```
#### 5.2.2 ConnectorFactory类
```python
class ConnectorFactory:
"""连接器工厂，根据设备和凭据创建合适的连接器"""
@staticmethod
def create_connector(device: NetworkDevice, credential: DeviceCredential) -> DeviceConnector:
"""创建设备连接器，目前主要使用ScrapliConnector"""
return ScrapliConnector(device, credential)
```
#### 5.2.3 DeviceConnector抽象基类
```python
class DeviceConnector(ABC):
"""设备连接器抽象基类"""
def __init__(self, device: NetworkDevice, credential: DeviceCredential):
self.device = device
self.credential = credential
self.connection = None
self.connected = False
self.last_activity = None
@abstractmethod
async def connect(self) -> bool:
"""连接到设备"""
@abstractmethod
async def disconnect(self) -> bool:
"""断开与设备的连接"""
@abstractmethod
async def send_command(self, command: str, timeout: int = 30) -> CommandResult:
"""发送命令到设备"""
@abstractmethod
async def send_commands(self, commands: List[str], timeout: int = 30) -> List[CommandResult]:
"""发送多个命令到设备"""
```
#### 5.2.4 ScrapliConnector实现
```python
class ScrapliConnector(DeviceConnector):
"""Scrapli连接器实现，支持混合异步/同步模式"""
async def connect(self) -> bool:
"""通过Scrapli连接到设备，支持多种平台"""
# 根据设备平台选择合适的驱动
# 支持异步和同步两种模式
# 实现自动重连机制
async def send_command(self, command: str, timeout: int = 30) -> CommandResult:
"""发送命令，支持异步/同步混合模式"""
# 根据连接对象类型选择调用方式
if asyncio.iscoroutinefunction(self.connection.send_command):
# 异步方法直接调用
response = await self.connection.send_command(command=command, timeout_ops=timeout)
else:
# 同步方法需要在线程中运行
response = await asyncio.to_thread(self.connection.send_command, command=command, timeout_ops=timeout)
```
#### 5.2.5 CommandResult类
```python
class CommandResult:
"""命令执行结果类"""
def __init__(self,
command: str,
output: str,
success: bool = True,
error_message: Optional[str] = None):
self.command = command
self.output = output
self.success = success
self.error_message = error_message
self.execution_time = datetime.now()
def to_dict(self) -> Dict[str, Any]:
"""转换为字典表示"""
```
## 6. 模板系统
### 6.1 模块概述
模板系统负责管理和渲染各种提示和命令模板，提高系统的灵活性和扩展性。系统实现了13个提示模板。
**实现文件**: `template_system.py`, `device_prompts.py`
### 6.2 核心类与接口
#### 6.2.1 TemplateManager类
```python
class TemplateManager:
"""高级提示模板管理器"""
def __init__(self):
self.templates = {}
self.templates_dir = os.path.join(os.getcwd(), "templates")
def register_template(self, name: str = None, description: str = None):
"""注册模板装饰器"""
def get_template(self, name: str) -> Optional[PromptTemplate]:
"""获取模板"""
def render_template(self, name: str, arguments: Dict[str, Any] = None) -> Optional[Union[str, List[Message]]]:
"""渲染模板"""
def list_templates(self) -> List[Dict[str, Any]]:
"""列出所有可用模板"""
```
#### 6.2.2 消息模型
```python
@dataclass
class Message:
"""基础消息类"""
role: str
content: str
class SystemMessage(Message):
"""系统消息"""
def __init__(self, content: str):
super().__init__(role="system", content=content)
class UserMessage(Message):
"""用户消息"""
def __init__(self, content: str):
super().__init__(role="user", content=content)
class AssistantMessage(Message):
"""助手消息"""
def __init__(self, content: str):
super().__init__(role="assistant", content=content)
```
### 6.3 实际模板实现
系统中实际实现的13个提示模板：
#### template_system.py中的模板 (8个)
1. `network_diagnosis` - 网络诊断提示模板
2. `device_configuration` - 设备配置会话模板
3. `device_diagnosis` - 网络设备诊断模板
4. `config_review` - 网络设备配置审查模板
5. `route_analysis` - 网络路由分析模板
6. `vlan_config` - VLAN配置生成模板
7. `network_topology` - 网络拓扑分析模板
8. `network_security` - 网络安全评估模板
#### device_prompts.py中的模板 (5个)
1. `cisco_device_analysis` - 思科设备分析模板
2. `huawei_device_analysis` - 华为设备分析模板
3. `network_troubleshooting` - 网络故障排除模板
4. `security_audit` - 网络安全审计模板
5. `performance_optimization` - 网络性能优化模板
#### 6.2.3 模板渲染函数
```python
async def render_template_with_resources(
template_name: str,
context: Dict[str, Any],
resource_uris: Dict[str, str],
resource_manager=None
) -> Union[str, List[Message]]:
"""渲染包含资源引用的模板"""
```
## 7. 扩展模块
### 7.1 拓扑发现模块
**实现文件**: `topology_discovery_improved.py`
该模块实现了基于CDP/LLDP协议的网络拓扑自动发现功能，支持多厂商设备的邻居发现和拓扑构建。
### 7.2 网络扫描模块
**实现文件**: `network_scanner.py`
该模块实现了网络范围扫描和设备发现功能，支持ping检测、端口扫描、SNMP检测等多种发现方式。
## 8. 数据持久化
系统采用JSON文件存储方式进行数据持久化：
- **设备数据**: `data/devices.json`
- **凭据数据**: `data/credentials.json`
- **资源缓存**: `resource_cache/` 目录
- **模板缓存**: `templates/` 目录
## 9. 模块间交互
```
MCP服务器 (server.py)
↓
├── 工具管理 (tool_manager.py) → 功能分类
├── 资源管理 (mcp_resources.py) → 数据提供
├── 模板系统 (template_system.py + device_prompts.py) → 提示生成
├── 设备管理 (network_devices.py) → 设备CRUD
├── 设备连接 (device_connector.py) → 连接管理
├── 拓扑发现 (topology_discovery_improved.py) → 拓扑构建
└── 网络扫描 (network_scanner.py) → 设备发现
```
这种模块化设计确保了系统的可维护性和扩展性，每个模块都有明确的职责和接口，便于独立开发和测试。