414 lines
14 KiB
Python
414 lines
14 KiB
Python
"""
|
||
Recon Agent (信息收集层) - LLM 驱动版
|
||
|
||
LLM 是真正的大脑！
|
||
- LLM 决定收集什么信息
|
||
- LLM 决定使用哪个工具
|
||
- LLM 决定何时信息足够
|
||
- LLM 动态调整收集策略
|
||
|
||
类型: ReAct (真正的!)
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
import re
|
||
from typing import List, Dict, Any, Optional
|
||
from dataclasses import dataclass
|
||
|
||
from .base import BaseAgent, AgentConfig, AgentResult, AgentType, AgentPattern
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
RECON_SYSTEM_PROMPT = """你是 DeepAudit 的信息收集 Agent，负责在安全审计前**自主**收集项目信息。
|
||
|
||
## 你的角色
|
||
你是信息收集的**大脑**，不是机械执行者。你需要：
|
||
1. 自主思考需要收集什么信息
|
||
2. 选择合适的工具获取信息
|
||
3. 根据发现动态调整策略
|
||
4. 判断何时信息收集足够
|
||
|
||
## 你可以使用的工具
|
||
|
||
### 文件系统
|
||
- **list_files**: 列出目录内容
|
||
参数: directory (str), recursive (bool), pattern (str), max_files (int)
|
||
|
||
- **read_file**: 读取文件内容
|
||
参数: file_path (str), start_line (int), end_line (int), max_lines (int)
|
||
|
||
- **search_code**: 代码关键字搜索
|
||
参数: keyword (str), max_results (int)
|
||
|
||
### 安全扫描
|
||
- **semgrep_scan**: Semgrep 静态分析扫描
|
||
- **npm_audit**: npm 依赖漏洞审计
|
||
- **safety_scan**: Python 依赖漏洞审计
|
||
- **gitleaks_scan**: 密钥/敏感信息泄露扫描
|
||
- **osv_scan**: OSV 通用依赖漏洞扫描
|
||
|
||
## 工作方式
|
||
每一步，你需要输出：
|
||
|
||
```
|
||
Thought: [分析当前状态，思考还需要什么信息]
|
||
Action: [工具名称]
|
||
Action Input: [JSON 格式的参数]
|
||
```
|
||
|
||
当你认为信息收集足够时，输出：
|
||
|
||
```
|
||
Thought: [总结收集到的信息]
|
||
Final Answer: [JSON 格式的收集结果]
|
||
```
|
||
|
||
## Final Answer 格式
|
||
```json
|
||
{
|
||
"project_structure": {
|
||
"directories": [],
|
||
"config_files": [],
|
||
"total_files": 数量
|
||
},
|
||
"tech_stack": {
|
||
"languages": [],
|
||
"frameworks": [],
|
||
"databases": []
|
||
},
|
||
"entry_points": [
|
||
{"type": "描述", "file": "路径", "line": 行号}
|
||
],
|
||
"high_risk_areas": ["路径列表"],
|
||
"dependencies": {},
|
||
"initial_findings": []
|
||
}
|
||
```
|
||
|
||
## 信息收集策略建议
|
||
1. 先 list_files 了解项目结构
|
||
2. 读取配置文件 (package.json, requirements.txt, go.mod 等) 识别技术栈
|
||
3. 搜索入口点模式 (routes, controllers, handlers)
|
||
4. 运行安全扫描发现初步问题
|
||
5. 根据发现继续深入
|
||
|
||
## 重要原则
|
||
1. **你是大脑** - 每一步都要思考，不要机械执行
|
||
2. **动态调整** - 根据发现调整策略
|
||
3. **效率优先** - 不要重复收集已有信息
|
||
4. **主动探索** - 发现有趣的东西要深入
|
||
|
||
现在开始收集项目信息！"""
|
||
|
||
|
||
@dataclass
|
||
class ReconStep:
|
||
"""信息收集步骤"""
|
||
thought: str
|
||
action: Optional[str] = None
|
||
action_input: Optional[Dict] = None
|
||
observation: Optional[str] = None
|
||
is_final: bool = False
|
||
final_answer: Optional[Dict] = None
|
||
|
||
|
||
class ReconAgent(BaseAgent):
|
||
"""
|
||
信息收集 Agent - LLM 驱动版
|
||
|
||
LLM 全程参与，自主决定：
|
||
1. 收集什么信息
|
||
2. 使用什么工具
|
||
3. 何时足够
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
llm_service,
|
||
tools: Dict[str, Any],
|
||
event_emitter=None,
|
||
):
|
||
config = AgentConfig(
|
||
name="Recon",
|
||
agent_type=AgentType.RECON,
|
||
pattern=AgentPattern.REACT,
|
||
max_iterations=15,
|
||
system_prompt=RECON_SYSTEM_PROMPT,
|
||
)
|
||
super().__init__(config, llm_service, tools, event_emitter)
|
||
|
||
self._conversation_history: List[Dict[str, str]] = []
|
||
self._steps: List[ReconStep] = []
|
||
|
||
def _get_tools_description(self) -> str:
|
||
"""生成工具描述"""
|
||
tools_info = []
|
||
for name, tool in self.tools.items():
|
||
if name.startswith("_"):
|
||
continue
|
||
desc = f"- {name}: {getattr(tool, 'description', 'No description')}"
|
||
tools_info.append(desc)
|
||
return "\n".join(tools_info)
|
||
|
||
def _parse_llm_response(self, response: str) -> ReconStep:
|
||
"""解析 LLM 响应"""
|
||
step = ReconStep(thought="")
|
||
|
||
# 提取 Thought
|
||
thought_match = re.search(r'Thought:\s*(.*?)(?=Action:|Final Answer:|$)', response, re.DOTALL)
|
||
if thought_match:
|
||
step.thought = thought_match.group(1).strip()
|
||
|
||
# 检查是否是最终答案
|
||
final_match = re.search(r'Final Answer:\s*(.*?)$', response, re.DOTALL)
|
||
if final_match:
|
||
step.is_final = True
|
||
try:
|
||
answer_text = final_match.group(1).strip()
|
||
answer_text = re.sub(r'```json\s*', '', answer_text)
|
||
answer_text = re.sub(r'```\s*', '', answer_text)
|
||
step.final_answer = json.loads(answer_text)
|
||
except json.JSONDecodeError:
|
||
step.final_answer = {"raw_answer": final_match.group(1).strip()}
|
||
return step
|
||
|
||
# 提取 Action
|
||
action_match = re.search(r'Action:\s*(\w+)', response)
|
||
if action_match:
|
||
step.action = action_match.group(1).strip()
|
||
|
||
# 提取 Action Input
|
||
input_match = re.search(r'Action Input:\s*(.*?)(?=Thought:|Action:|Observation:|$)', response, re.DOTALL)
|
||
if input_match:
|
||
input_text = input_match.group(1).strip()
|
||
input_text = re.sub(r'```json\s*', '', input_text)
|
||
input_text = re.sub(r'```\s*', '', input_text)
|
||
try:
|
||
step.action_input = json.loads(input_text)
|
||
except json.JSONDecodeError:
|
||
step.action_input = {"raw_input": input_text}
|
||
|
||
return step
|
||
|
||
async def _execute_tool(self, tool_name: str, tool_input: Dict) -> str:
|
||
"""执行工具"""
|
||
tool = self.tools.get(tool_name)
|
||
|
||
if not tool:
|
||
return f"错误: 工具 '{tool_name}' 不存在。可用工具: {list(self.tools.keys())}"
|
||
|
||
try:
|
||
self._tool_calls += 1
|
||
await self.emit_tool_call(tool_name, tool_input)
|
||
|
||
import time
|
||
start = time.time()
|
||
|
||
result = await tool.execute(**tool_input)
|
||
|
||
duration_ms = int((time.time() - start) * 1000)
|
||
await self.emit_tool_result(tool_name, str(result.data)[:200], duration_ms)
|
||
|
||
if result.success:
|
||
output = str(result.data)
|
||
if len(output) > 4000:
|
||
output = output[:4000] + f"\n\n... [输出已截断，共 {len(str(result.data))} 字符]"
|
||
return output
|
||
else:
|
||
return f"工具执行失败: {result.error}"
|
||
|
||
except Exception as e:
|
||
logger.error(f"Tool execution error: {e}")
|
||
return f"工具执行错误: {str(e)}"
|
||
|
||
async def run(self, input_data: Dict[str, Any]) -> AgentResult:
|
||
"""
|
||
执行信息收集 - LLM 全程参与！
|
||
"""
|
||
import time
|
||
start_time = time.time()
|
||
|
||
project_info = input_data.get("project_info", {})
|
||
config = input_data.get("config", {})
|
||
task = input_data.get("task", "")
|
||
task_context = input_data.get("task_context", "")
|
||
|
||
# 构建初始消息
|
||
initial_message = f"""请开始收集项目信息。
|
||
|
||
## 项目基本信息
|
||
- 名称: {project_info.get('name', 'unknown')}
|
||
- 根目录: {project_info.get('root', '.')}
|
||
|
||
## 任务上下文
|
||
{task_context or task or '进行全面的信息收集，为安全审计做准备。'}
|
||
|
||
## 可用工具
|
||
{self._get_tools_description()}
|
||
|
||
请开始你的信息收集工作。首先思考应该收集什么信息，然后选择合适的工具。"""
|
||
|
||
# 初始化对话历史
|
||
self._conversation_history = [
|
||
{"role": "system", "content": self.config.system_prompt},
|
||
{"role": "user", "content": initial_message},
|
||
]
|
||
|
||
self._steps = []
|
||
final_result = None
|
||
|
||
await self.emit_thinking("🔍 Recon Agent 启动，LLM 开始自主收集信息...")
|
||
|
||
try:
|
||
for iteration in range(self.config.max_iterations):
|
||
if self.is_cancelled:
|
||
break
|
||
|
||
self._iteration = iteration + 1
|
||
|
||
# 🔥 发射 LLM 开始思考事件
|
||
await self.emit_llm_start(iteration + 1)
|
||
|
||
# 🔥 调用 LLM 进行思考和决策
|
||
response = await self.llm_service.chat_completion_raw(
|
||
messages=self._conversation_history,
|
||
temperature=0.1,
|
||
max_tokens=2048,
|
||
)
|
||
|
||
llm_output = response.get("content", "")
|
||
tokens_this_round = response.get("usage", {}).get("total_tokens", 0)
|
||
self._total_tokens += tokens_this_round
|
||
|
||
# 解析 LLM 响应
|
||
step = self._parse_llm_response(llm_output)
|
||
self._steps.append(step)
|
||
|
||
# 🔥 发射 LLM 思考内容事件 - 展示 LLM 在想什么
|
||
if step.thought:
|
||
await self.emit_llm_thought(step.thought, iteration + 1)
|
||
|
||
# 添加 LLM 响应到历史
|
||
self._conversation_history.append({
|
||
"role": "assistant",
|
||
"content": llm_output,
|
||
})
|
||
|
||
# 检查是否完成
|
||
if step.is_final:
|
||
await self.emit_llm_decision("完成信息收集", "LLM 判断已收集足够信息")
|
||
await self.emit_llm_complete(
|
||
f"信息收集完成，共 {self._iteration} 轮思考",
|
||
self._total_tokens
|
||
)
|
||
final_result = step.final_answer
|
||
break
|
||
|
||
# 执行工具
|
||
if step.action:
|
||
# 🔥 发射 LLM 动作决策事件
|
||
await self.emit_llm_action(step.action, step.action_input or {})
|
||
|
||
observation = await self._execute_tool(
|
||
step.action,
|
||
step.action_input or {}
|
||
)
|
||
|
||
step.observation = observation
|
||
|
||
# 🔥 发射 LLM 观察事件
|
||
await self.emit_llm_observation(observation)
|
||
|
||
# 添加观察结果到历史
|
||
self._conversation_history.append({
|
||
"role": "user",
|
||
"content": f"Observation:\n{observation}",
|
||
})
|
||
else:
|
||
# LLM 没有选择工具，提示它继续
|
||
await self.emit_llm_decision("继续思考", "LLM 需要更多信息")
|
||
self._conversation_history.append({
|
||
"role": "user",
|
||
"content": "请继续，选择一个工具执行，或者如果信息收集完成，输出 Final Answer。",
|
||
})
|
||
|
||
# 处理结果
|
||
duration_ms = int((time.time() - start_time) * 1000)
|
||
|
||
# 如果没有最终结果，从历史中汇总
|
||
if not final_result:
|
||
final_result = self._summarize_from_steps()
|
||
|
||
await self.emit_event(
|
||
"info",
|
||
f"🎯 Recon Agent 完成: {self._iteration} 轮迭代, {self._tool_calls} 次工具调用"
|
||
)
|
||
|
||
return AgentResult(
|
||
success=True,
|
||
data=final_result,
|
||
iterations=self._iteration,
|
||
tool_calls=self._tool_calls,
|
||
tokens_used=self._total_tokens,
|
||
duration_ms=duration_ms,
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Recon Agent failed: {e}", exc_info=True)
|
||
return AgentResult(success=False, error=str(e))
|
||
|
||
def _summarize_from_steps(self) -> Dict[str, Any]:
|
||
"""从步骤中汇总结果"""
|
||
# 默认结果结构
|
||
result = {
|
||
"project_structure": {},
|
||
"tech_stack": {
|
||
"languages": [],
|
||
"frameworks": [],
|
||
"databases": [],
|
||
},
|
||
"entry_points": [],
|
||
"high_risk_areas": [],
|
||
"dependencies": {},
|
||
"initial_findings": [],
|
||
}
|
||
|
||
# 从步骤的观察结果中提取信息
|
||
for step in self._steps:
|
||
if step.observation:
|
||
# 尝试从观察中识别技术栈等信息
|
||
obs_lower = step.observation.lower()
|
||
|
||
if "package.json" in obs_lower:
|
||
result["tech_stack"]["languages"].append("JavaScript/TypeScript")
|
||
if "requirements.txt" in obs_lower or "setup.py" in obs_lower:
|
||
result["tech_stack"]["languages"].append("Python")
|
||
if "go.mod" in obs_lower:
|
||
result["tech_stack"]["languages"].append("Go")
|
||
|
||
# 识别框架
|
||
if "react" in obs_lower:
|
||
result["tech_stack"]["frameworks"].append("React")
|
||
if "django" in obs_lower:
|
||
result["tech_stack"]["frameworks"].append("Django")
|
||
if "fastapi" in obs_lower:
|
||
result["tech_stack"]["frameworks"].append("FastAPI")
|
||
if "express" in obs_lower:
|
||
result["tech_stack"]["frameworks"].append("Express")
|
||
|
||
# 去重
|
||
result["tech_stack"]["languages"] = list(set(result["tech_stack"]["languages"]))
|
||
result["tech_stack"]["frameworks"] = list(set(result["tech_stack"]["frameworks"]))
|
||
|
||
return result
|
||
|
||
def get_conversation_history(self) -> List[Dict[str, str]]:
|
||
"""获取对话历史"""
|
||
return self._conversation_history
|
||
|
||
def get_steps(self) -> List[ReconStep]:
|
||
"""获取执行步骤"""
|
||
return self._steps
|