306 lines
10 KiB
Python
306 lines
10 KiB
Python
|
"""
|
|||
|
|
扫描完成工具
|
|||
|
|
|
|||
|
|
用于主Agent结束安全审计任务，确保所有子Agent已完成。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
from datetime import datetime, timezone
|
|||
|
|
from typing import Optional, List, Dict, Any
|
|||
|
|
from pydantic import BaseModel, Field
|
|||
|
|
|
|||
|
|
from .base import AgentTool, ToolResult
|
|||
|
|
from ..core.registry import agent_registry
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FinishScanInput(BaseModel):
|
|||
|
|
"""扫描完成输入参数"""
|
|||
|
|
content: str = Field(
|
|||
|
|
...,
|
|||
|
|
description="最终扫描报告内容，包含所有发现的漏洞总结"
|
|||
|
|
)
|
|||
|
|
success: bool = Field(
|
|||
|
|
default=True,
|
|||
|
|
description="扫描是否成功完成"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FinishScanTool(AgentTool):
|
|||
|
|
"""
|
|||
|
|
扫描完成工具
|
|||
|
|
|
|||
|
|
只有根Agent（主Agent）才能使用此工具来正式结束安全审计任务。
|
|||
|
|
|
|||
|
|
使用前置条件：
|
|||
|
|
1. 所有子Agent必须已完成（completed, failed, 或 stopped 状态）
|
|||
|
|
2. 必须提供最终报告内容
|
|||
|
|
|
|||
|
|
使用约束：
|
|||
|
|
- 子Agent必须使用 agent_finish 工具
|
|||
|
|
- 根Agent必须使用此工具
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, agent_id: str, agent_state=None, tracer=None):
|
|||
|
|
super().__init__()
|
|||
|
|
self.agent_id = agent_id
|
|||
|
|
self.agent_state = agent_state
|
|||
|
|
self.tracer = tracer
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def name(self) -> str:
|
|||
|
|
return "finish_scan"
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def description(self) -> str:
|
|||
|
|
return """完成整个安全扫描并生成最终报告。
|
|||
|
|
|
|||
|
|
只有根Agent（主编排Agent）才能使用此工具。
|
|||
|
|
|
|||
|
|
使用条件：
|
|||
|
|
1. 所有子Agent必须已完成
|
|||
|
|
2. 必须提供完整的扫描总结
|
|||
|
|
|
|||
|
|
参数:
|
|||
|
|
- content: 最终扫描报告内容，包含：
|
|||
|
|
- 扫描概述
|
|||
|
|
- 发现的漏洞列表
|
|||
|
|
- 风险评估
|
|||
|
|
- 修复建议
|
|||
|
|
- success: 扫描是否成功完成
|
|||
|
|
|
|||
|
|
不要使用此工具:
|
|||
|
|
- 如果你是子Agent（请使用 agent_finish）
|
|||
|
|
- 如果还有子Agent正在运行"""
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def args_schema(self):
|
|||
|
|
return FinishScanInput
|
|||
|
|
|
|||
|
|
async def _execute(
|
|||
|
|
self,
|
|||
|
|
content: str,
|
|||
|
|
success: bool = True,
|
|||
|
|
**kwargs
|
|||
|
|
) -> ToolResult:
|
|||
|
|
"""执行扫描完成"""
|
|||
|
|
|
|||
|
|
# 验证是否为根Agent
|
|||
|
|
validation_error = self._validate_root_agent()
|
|||
|
|
if validation_error:
|
|||
|
|
return validation_error
|
|||
|
|
|
|||
|
|
# 验证内容
|
|||
|
|
if not content or not content.strip():
|
|||
|
|
return ToolResult(success=False, error="报告内容不能为空")
|
|||
|
|
|
|||
|
|
# 检查是否有活跃的子Agent
|
|||
|
|
active_check = self._check_active_agents()
|
|||
|
|
if active_check:
|
|||
|
|
return active_check
|
|||
|
|
|
|||
|
|
# 更新状态
|
|||
|
|
final_result = {
|
|||
|
|
"scan_completed": True,
|
|||
|
|
"content": content.strip(),
|
|||
|
|
"success": success,
|
|||
|
|
"completed_at": datetime.now(timezone.utc).isoformat(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 收集所有发现
|
|||
|
|
all_findings = self._collect_all_findings()
|
|||
|
|
final_result["total_findings"] = len(all_findings)
|
|||
|
|
final_result["findings_summary"] = self._summarize_findings(all_findings)
|
|||
|
|
|
|||
|
|
# 更新Agent状态
|
|||
|
|
if self.agent_state:
|
|||
|
|
self.agent_state.set_completed(final_result)
|
|||
|
|
|
|||
|
|
agent_registry.update_agent_status(
|
|||
|
|
self.agent_id,
|
|||
|
|
"completed" if success else "failed",
|
|||
|
|
final_result,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 保存到追踪器
|
|||
|
|
if self.tracer:
|
|||
|
|
try:
|
|||
|
|
self.tracer.set_final_scan_result(
|
|||
|
|
content=content.strip(),
|
|||
|
|
success=success,
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"Failed to update tracer: {e}")
|
|||
|
|
|
|||
|
|
# 获取统计信息
|
|||
|
|
stats = agent_registry.get_statistics()
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
data={
|
|||
|
|
"scan_completed": True,
|
|||
|
|
"message": "扫描已成功完成" if success else "扫描完成但有错误",
|
|||
|
|
"report_length": len(content),
|
|||
|
|
"total_findings": len(all_findings),
|
|||
|
|
"agent_statistics": stats,
|
|||
|
|
},
|
|||
|
|
metadata=final_result,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def _validate_root_agent(self) -> Optional[ToolResult]:
|
|||
|
|
"""验证是否为根Agent"""
|
|||
|
|
# 检查是否有父Agent
|
|||
|
|
parent_id = agent_registry.get_parent(self.agent_id)
|
|||
|
|
|
|||
|
|
if parent_id is not None:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="此工具只能由根Agent使用。子Agent请使用 agent_finish 工具。"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 检查是否为注册的根Agent
|
|||
|
|
root_id = agent_registry.get_root_agent_id()
|
|||
|
|
if root_id and root_id != self.agent_id:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error=f"当前Agent不是根Agent。根Agent ID: {root_id}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def _check_active_agents(self) -> Optional[ToolResult]:
|
|||
|
|
"""检查是否有活跃的子Agent"""
|
|||
|
|
try:
|
|||
|
|
tree = agent_registry.get_agent_tree()
|
|||
|
|
|
|||
|
|
running_agents = []
|
|||
|
|
waiting_agents = []
|
|||
|
|
stopping_agents = []
|
|||
|
|
|
|||
|
|
for agent_id, node in tree["nodes"].items():
|
|||
|
|
# 跳过当前Agent
|
|||
|
|
if agent_id == self.agent_id:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
status = node.get("status", "")
|
|||
|
|
|
|||
|
|
if status == "running":
|
|||
|
|
running_agents.append({
|
|||
|
|
"id": agent_id,
|
|||
|
|
"name": node.get("name", "Unknown"),
|
|||
|
|
"task": node.get("task", "No task description")[:80],
|
|||
|
|
})
|
|||
|
|
elif status == "waiting":
|
|||
|
|
waiting_agents.append({
|
|||
|
|
"id": agent_id,
|
|||
|
|
"name": node.get("name", "Unknown"),
|
|||
|
|
})
|
|||
|
|
elif status == "stopping":
|
|||
|
|
stopping_agents.append({
|
|||
|
|
"id": agent_id,
|
|||
|
|
"name": node.get("name", "Unknown"),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
if running_agents or stopping_agents:
|
|||
|
|
message_parts = ["无法完成扫描，还有活跃的子Agent:"]
|
|||
|
|
|
|||
|
|
if running_agents:
|
|||
|
|
message_parts.append("\n\n运行中的Agent:")
|
|||
|
|
for agent in running_agents:
|
|||
|
|
message_parts.append(
|
|||
|
|
f" - {agent['name']} ({agent['id']}): {agent['task']}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if stopping_agents:
|
|||
|
|
message_parts.append("\n\n正在停止的Agent:")
|
|||
|
|
for agent in stopping_agents:
|
|||
|
|
message_parts.append(f" - {agent['name']} ({agent['id']})")
|
|||
|
|
|
|||
|
|
message_parts.extend([
|
|||
|
|
"\n\n建议操作:",
|
|||
|
|
"1. 使用 wait_for_message 等待所有Agent完成",
|
|||
|
|
"2. 使用 view_agent_graph 查看Agent状态",
|
|||
|
|
"3. 如果需要紧急结束，发送消息要求Agent完成",
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error="\n".join(message_parts),
|
|||
|
|
metadata={
|
|||
|
|
"running_count": len(running_agents),
|
|||
|
|
"stopping_count": len(stopping_agents),
|
|||
|
|
"waiting_count": len(waiting_agents),
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"Failed to check active agents: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def _collect_all_findings(self) -> List[Dict[str, Any]]:
|
|||
|
|
"""收集所有子Agent的发现"""
|
|||
|
|
all_findings = []
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
tree = agent_registry.get_agent_tree()
|
|||
|
|
|
|||
|
|
for agent_id, node in tree["nodes"].items():
|
|||
|
|
result = node.get("result")
|
|||
|
|
if not result:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 从result中提取findings
|
|||
|
|
if isinstance(result, dict):
|
|||
|
|
findings = result.get("findings", [])
|
|||
|
|
if isinstance(findings, list):
|
|||
|
|
for finding in findings:
|
|||
|
|
if isinstance(finding, dict):
|
|||
|
|
finding["discovered_by"] = {
|
|||
|
|
"agent_id": agent_id,
|
|||
|
|
"agent_name": node.get("name", "Unknown"),
|
|||
|
|
}
|
|||
|
|
all_findings.append(finding)
|
|||
|
|
elif isinstance(finding, str):
|
|||
|
|
all_findings.append({
|
|||
|
|
"description": finding,
|
|||
|
|
"discovered_by": {
|
|||
|
|
"agent_id": agent_id,
|
|||
|
|
"agent_name": node.get("name", "Unknown"),
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"Failed to collect findings: {e}")
|
|||
|
|
|
|||
|
|
return all_findings
|
|||
|
|
|
|||
|
|
def _summarize_findings(self, findings: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|||
|
|
"""生成发现摘要"""
|
|||
|
|
severity_counts = {
|
|||
|
|
"critical": 0,
|
|||
|
|
"high": 0,
|
|||
|
|
"medium": 0,
|
|||
|
|
"low": 0,
|
|||
|
|
"info": 0,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type_counts = {}
|
|||
|
|
|
|||
|
|
for finding in findings:
|
|||
|
|
# 统计严重性
|
|||
|
|
severity = finding.get("severity", "medium").lower()
|
|||
|
|
if severity in severity_counts:
|
|||
|
|
severity_counts[severity] += 1
|
|||
|
|
|
|||
|
|
# 统计类型
|
|||
|
|
vuln_type = finding.get("vulnerability_type", finding.get("type", "unknown"))
|
|||
|
|
type_counts[vuln_type] = type_counts.get(vuln_type, 0) + 1
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"total": len(findings),
|
|||
|
|
"by_severity": severity_counts,
|
|||
|
|
"by_type": type_counts,
|
|||
|
|
}
|