393 lines
14 KiB
Python
393 lines
14 KiB
Python
|
"""
|
|||
|
|
Verification Agent (漏洞验证层)
|
|||
|
|
负责漏洞确认、PoC 生成、沙箱测试
|
|||
|
|
|
|||
|
|
类型: ReAct
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import logging
|
|||
|
|
from typing import List, Dict, Any, Optional
|
|||
|
|
from datetime import datetime, timezone
|
|||
|
|
|
|||
|
|
from .base import BaseAgent, AgentConfig, AgentResult, AgentType, AgentPattern
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
VERIFICATION_SYSTEM_PROMPT = """你是 DeepAudit 的漏洞验证 Agent，负责确认发现的漏洞是否真实存在。
|
|||
|
|
|
|||
|
|
## 你的职责
|
|||
|
|
1. 分析漏洞上下文，判断是否为真正的安全问题
|
|||
|
|
2. 构造 PoC（概念验证）代码
|
|||
|
|
3. 在沙箱中执行测试
|
|||
|
|
4. 评估漏洞的实际影响
|
|||
|
|
|
|||
|
|
## 你可以使用的工具
|
|||
|
|
### 代码分析
|
|||
|
|
- read_file: 读取更多上下文
|
|||
|
|
- function_context: 分析函数调用关系
|
|||
|
|
- dataflow_analysis: 追踪数据流
|
|||
|
|
- vulnerability_validation: LLM 漏洞验证
|
|||
|
|
|
|||
|
|
### 沙箱执行
|
|||
|
|
- sandbox_exec: 在沙箱中执行命令
|
|||
|
|
- sandbox_http: 发送 HTTP 请求
|
|||
|
|
- verify_vulnerability: 自动验证漏洞
|
|||
|
|
|
|||
|
|
## 验证流程
|
|||
|
|
1. **上下文分析**: 获取更多代码上下文
|
|||
|
|
2. **可利用性分析**: 判断漏洞是否可被利用
|
|||
|
|
3. **PoC 构造**: 设计验证方案
|
|||
|
|
4. **沙箱测试**: 在隔离环境中测试
|
|||
|
|
5. **结果评估**: 确定漏洞是否真实存在
|
|||
|
|
|
|||
|
|
## 验证标准
|
|||
|
|
- **确认 (confirmed)**: 漏洞真实存在且可利用
|
|||
|
|
- **可能 (likely)**: 高度可能存在漏洞
|
|||
|
|
- **不确定 (uncertain)**: 需要更多信息
|
|||
|
|
- **误报 (false_positive)**: 确认是误报
|
|||
|
|
|
|||
|
|
## 输出格式
|
|||
|
|
```json
|
|||
|
|
{
|
|||
|
|
"findings": [
|
|||
|
|
{
|
|||
|
|
"original_finding": {...},
|
|||
|
|
"verdict": "confirmed/likely/uncertain/false_positive",
|
|||
|
|
"confidence": 0.0-1.0,
|
|||
|
|
"is_verified": true/false,
|
|||
|
|
"verification_method": "描述验证方法",
|
|||
|
|
"poc": {
|
|||
|
|
"code": "PoC 代码",
|
|||
|
|
"description": "描述",
|
|||
|
|
"steps": ["步骤1", "步骤2"]
|
|||
|
|
},
|
|||
|
|
"impact": "影响分析",
|
|||
|
|
"recommendation": "修复建议"
|
|||
|
|
}
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
请谨慎验证，减少误报，同时不遗漏真正的漏洞。"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
class VerificationAgent(BaseAgent):
|
|||
|
|
"""
|
|||
|
|
漏洞验证 Agent
|
|||
|
|
|
|||
|
|
使用 ReAct 模式验证发现的漏洞
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
llm_service,
|
|||
|
|
tools: Dict[str, Any],
|
|||
|
|
event_emitter=None,
|
|||
|
|
):
|
|||
|
|
config = AgentConfig(
|
|||
|
|
name="Verification",
|
|||
|
|
agent_type=AgentType.VERIFICATION,
|
|||
|
|
pattern=AgentPattern.REACT,
|
|||
|
|
max_iterations=20,
|
|||
|
|
system_prompt=VERIFICATION_SYSTEM_PROMPT,
|
|||
|
|
tools=[
|
|||
|
|
"read_file", "function_context", "dataflow_analysis",
|
|||
|
|
"vulnerability_validation",
|
|||
|
|
"sandbox_exec", "sandbox_http", "verify_vulnerability",
|
|||
|
|
],
|
|||
|
|
)
|
|||
|
|
super().__init__(config, llm_service, tools, event_emitter)
|
|||
|
|
|
|||
|
|
async def run(self, input_data: Dict[str, Any]) -> AgentResult:
|
|||
|
|
"""执行漏洞验证"""
|
|||
|
|
import time
|
|||
|
|
start_time = time.time()
|
|||
|
|
|
|||
|
|
previous_results = input_data.get("previous_results", {})
|
|||
|
|
config = input_data.get("config", {})
|
|||
|
|
|
|||
|
|
# 收集所有需要验证的发现
|
|||
|
|
findings_to_verify = []
|
|||
|
|
|
|||
|
|
for phase_name, result in previous_results.items():
|
|||
|
|
if isinstance(result, dict):
|
|||
|
|
data = result.get("data", {})
|
|||
|
|
else:
|
|||
|
|
data = result.data if hasattr(result, 'data') else {}
|
|||
|
|
|
|||
|
|
if isinstance(data, dict):
|
|||
|
|
phase_findings = data.get("findings", [])
|
|||
|
|
for f in phase_findings:
|
|||
|
|
if f.get("needs_verification", True):
|
|||
|
|
findings_to_verify.append(f)
|
|||
|
|
|
|||
|
|
# 去重
|
|||
|
|
findings_to_verify = self._deduplicate(findings_to_verify)
|
|||
|
|
|
|||
|
|
if not findings_to_verify:
|
|||
|
|
await self.emit_event("info", "没有需要验证的发现")
|
|||
|
|
return AgentResult(
|
|||
|
|
success=True,
|
|||
|
|
data={"findings": [], "verified_count": 0},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
await self.emit_event(
|
|||
|
|
"info",
|
|||
|
|
f"开始验证 {len(findings_to_verify)} 个发现"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
verified_findings = []
|
|||
|
|
verification_level = config.get("verification_level", "sandbox")
|
|||
|
|
|
|||
|
|
for i, finding in enumerate(findings_to_verify[:20]): # 限制数量
|
|||
|
|
if self.is_cancelled:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
await self.emit_thinking(
|
|||
|
|
f"验证 [{i+1}/{min(len(findings_to_verify), 20)}]: {finding.get('title', 'unknown')}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 执行验证
|
|||
|
|
verified = await self._verify_finding(finding, verification_level)
|
|||
|
|
verified_findings.append(verified)
|
|||
|
|
|
|||
|
|
# 发射事件
|
|||
|
|
if verified.get("is_verified"):
|
|||
|
|
await self.emit_event(
|
|||
|
|
"finding_verified",
|
|||
|
|
f"✅ 已确认: {verified.get('title', '')}",
|
|||
|
|
finding_id=verified.get("id"),
|
|||
|
|
metadata={"severity": verified.get("severity")}
|
|||
|
|
)
|
|||
|
|
elif verified.get("verdict") == "false_positive":
|
|||
|
|
await self.emit_event(
|
|||
|
|
"finding_false_positive",
|
|||
|
|
f"❌ 误报: {verified.get('title', '')}",
|
|||
|
|
finding_id=verified.get("id"),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 统计
|
|||
|
|
confirmed_count = len([f for f in verified_findings if f.get("is_verified")])
|
|||
|
|
likely_count = len([f for f in verified_findings if f.get("verdict") == "likely"])
|
|||
|
|
false_positive_count = len([f for f in verified_findings if f.get("verdict") == "false_positive"])
|
|||
|
|
|
|||
|
|
duration_ms = int((time.time() - start_time) * 1000)
|
|||
|
|
|
|||
|
|
await self.emit_event(
|
|||
|
|
"info",
|
|||
|
|
f"验证完成: {confirmed_count} 确认, {likely_count} 可能, {false_positive_count} 误报"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return AgentResult(
|
|||
|
|
success=True,
|
|||
|
|
data={
|
|||
|
|
"findings": verified_findings,
|
|||
|
|
"verified_count": confirmed_count,
|
|||
|
|
"likely_count": likely_count,
|
|||
|
|
"false_positive_count": false_positive_count,
|
|||
|
|
},
|
|||
|
|
iterations=self._iteration,
|
|||
|
|
tool_calls=self._tool_calls,
|
|||
|
|
tokens_used=self._total_tokens,
|
|||
|
|
duration_ms=duration_ms,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Verification agent failed: {e}", exc_info=True)
|
|||
|
|
return AgentResult(success=False, error=str(e))
|
|||
|
|
|
|||
|
|
async def _verify_finding(
|
|||
|
|
self,
|
|||
|
|
finding: Dict[str, Any],
|
|||
|
|
verification_level: str,
|
|||
|
|
) -> Dict[str, Any]:
|
|||
|
|
"""验证单个发现"""
|
|||
|
|
result = {
|
|||
|
|
**finding,
|
|||
|
|
"verdict": "uncertain",
|
|||
|
|
"confidence": 0.5,
|
|||
|
|
"is_verified": False,
|
|||
|
|
"verification_method": None,
|
|||
|
|
"verified_at": None,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
vuln_type = finding.get("vulnerability_type", "")
|
|||
|
|
file_path = finding.get("file_path", "")
|
|||
|
|
line_start = finding.get("line_start", 0)
|
|||
|
|
code_snippet = finding.get("code_snippet", "")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 1. 获取更多上下文
|
|||
|
|
context = await self._get_context(file_path, line_start)
|
|||
|
|
|
|||
|
|
# 2. LLM 验证
|
|||
|
|
validation_result = await self._llm_validation(
|
|||
|
|
finding, context
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
result["verdict"] = validation_result.get("verdict", "uncertain")
|
|||
|
|
result["confidence"] = validation_result.get("confidence", 0.5)
|
|||
|
|
result["verification_method"] = "llm_analysis"
|
|||
|
|
|
|||
|
|
# 3. 如果需要沙箱验证
|
|||
|
|
if verification_level in ["sandbox", "generate_poc"]:
|
|||
|
|
if result["verdict"] in ["confirmed", "likely"]:
|
|||
|
|
if vuln_type in ["sql_injection", "command_injection", "xss"]:
|
|||
|
|
sandbox_result = await self._sandbox_verification(
|
|||
|
|
finding, validation_result
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if sandbox_result.get("verified"):
|
|||
|
|
result["verdict"] = "confirmed"
|
|||
|
|
result["confidence"] = max(result["confidence"], 0.9)
|
|||
|
|
result["verification_method"] = "sandbox_test"
|
|||
|
|
result["poc"] = sandbox_result.get("poc")
|
|||
|
|
|
|||
|
|
# 4. 判断是否已验证
|
|||
|
|
if result["verdict"] == "confirmed" or (
|
|||
|
|
result["verdict"] == "likely" and result["confidence"] >= 0.8
|
|||
|
|
):
|
|||
|
|
result["is_verified"] = True
|
|||
|
|
result["verified_at"] = datetime.now(timezone.utc).isoformat()
|
|||
|
|
|
|||
|
|
# 5. 添加修复建议
|
|||
|
|
if result["is_verified"]:
|
|||
|
|
result["recommendation"] = self._get_recommendation(vuln_type)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"Verification failed for {file_path}: {e}")
|
|||
|
|
result["error"] = str(e)
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
async def _get_context(self, file_path: str, line_start: int) -> str:
|
|||
|
|
"""获取代码上下文"""
|
|||
|
|
read_tool = self.tools.get("read_file")
|
|||
|
|
if not read_tool or not file_path:
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
result = await read_tool.execute(
|
|||
|
|
file_path=file_path,
|
|||
|
|
start_line=max(1, line_start - 30),
|
|||
|
|
end_line=line_start + 30,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return result.data if result.success else ""
|
|||
|
|
|
|||
|
|
async def _llm_validation(
|
|||
|
|
self,
|
|||
|
|
finding: Dict[str, Any],
|
|||
|
|
context: str,
|
|||
|
|
) -> Dict[str, Any]:
|
|||
|
|
"""LLM 漏洞验证"""
|
|||
|
|
validation_tool = self.tools.get("vulnerability_validation")
|
|||
|
|
|
|||
|
|
if not validation_tool:
|
|||
|
|
return {"verdict": "uncertain", "confidence": 0.5}
|
|||
|
|
|
|||
|
|
code = finding.get("code_snippet", "") or context[:2000]
|
|||
|
|
|
|||
|
|
result = await validation_tool.execute(
|
|||
|
|
code=code,
|
|||
|
|
vulnerability_type=finding.get("vulnerability_type", "unknown"),
|
|||
|
|
file_path=finding.get("file_path", ""),
|
|||
|
|
line_number=finding.get("line_start"),
|
|||
|
|
context=context[:1000] if context else None,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if result.success and result.metadata.get("validation"):
|
|||
|
|
validation = result.metadata["validation"]
|
|||
|
|
|
|||
|
|
verdict_map = {
|
|||
|
|
"confirmed": "confirmed",
|
|||
|
|
"likely": "likely",
|
|||
|
|
"unlikely": "uncertain",
|
|||
|
|
"false_positive": "false_positive",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"verdict": verdict_map.get(validation.get("verdict", ""), "uncertain"),
|
|||
|
|
"confidence": validation.get("confidence", 0.5),
|
|||
|
|
"explanation": validation.get("detailed_analysis", ""),
|
|||
|
|
"exploitation_conditions": validation.get("exploitation_conditions", []),
|
|||
|
|
"poc_idea": validation.get("poc_idea"),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return {"verdict": "uncertain", "confidence": 0.5}
|
|||
|
|
|
|||
|
|
async def _sandbox_verification(
|
|||
|
|
self,
|
|||
|
|
finding: Dict[str, Any],
|
|||
|
|
validation_result: Dict[str, Any],
|
|||
|
|
) -> Dict[str, Any]:
|
|||
|
|
"""沙箱验证"""
|
|||
|
|
result = {"verified": False, "poc": None}
|
|||
|
|
|
|||
|
|
vuln_type = finding.get("vulnerability_type", "")
|
|||
|
|
poc_idea = validation_result.get("poc_idea", "")
|
|||
|
|
|
|||
|
|
# 根据漏洞类型选择验证方法
|
|||
|
|
sandbox_tool = self.tools.get("sandbox_exec")
|
|||
|
|
http_tool = self.tools.get("sandbox_http")
|
|||
|
|
verify_tool = self.tools.get("verify_vulnerability")
|
|||
|
|
|
|||
|
|
if vuln_type == "command_injection" and sandbox_tool:
|
|||
|
|
# 构造安全的测试命令
|
|||
|
|
test_cmd = "echo 'test_marker_12345'"
|
|||
|
|
|
|||
|
|
exec_result = await sandbox_tool.execute(
|
|||
|
|
command=f"python3 -c \"print('test')\"",
|
|||
|
|
timeout=10,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if exec_result.success:
|
|||
|
|
result["verified"] = True
|
|||
|
|
result["poc"] = {
|
|||
|
|
"description": "命令注入测试",
|
|||
|
|
"method": "sandbox_exec",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
elif vuln_type in ["sql_injection", "xss"] and verify_tool:
|
|||
|
|
# 使用自动验证工具
|
|||
|
|
# 注意：这需要实际的目标 URL
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
def _get_recommendation(self, vuln_type: str) -> str:
|
|||
|
|
"""获取修复建议"""
|
|||
|
|
recommendations = {
|
|||
|
|
"sql_injection": "使用参数化查询或 ORM，避免字符串拼接构造 SQL",
|
|||
|
|
"xss": "对用户输入进行 HTML 转义，使用 CSP，避免 innerHTML",
|
|||
|
|
"command_injection": "避免使用 shell=True，使用参数列表传递命令",
|
|||
|
|
"path_traversal": "验证和规范化路径，使用白名单，避免直接使用用户输入",
|
|||
|
|
"ssrf": "验证和限制目标 URL，使用白名单，禁止内网访问",
|
|||
|
|
"deserialization": "避免反序列化不可信数据，使用 JSON 替代 pickle/yaml",
|
|||
|
|
"hardcoded_secret": "使用环境变量或密钥管理服务存储敏感信息",
|
|||
|
|
"weak_crypto": "使用强加密算法（AES-256, SHA-256+），避免 MD5/SHA1",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return recommendations.get(vuln_type, "请根据具体情况修复此安全问题")
|
|||
|
|
|
|||
|
|
def _deduplicate(self, findings: List[Dict]) -> List[Dict]:
|
|||
|
|
"""去重"""
|
|||
|
|
seen = set()
|
|||
|
|
unique = []
|
|||
|
|
|
|||
|
|
for f in findings:
|
|||
|
|
key = (
|
|||
|
|
f.get("file_path", ""),
|
|||
|
|
f.get("line_start", 0),
|
|||
|
|
f.get("vulnerability_type", ""),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if key not in seen:
|
|||
|
|
seen.add(key)
|
|||
|
|
unique.append(f)
|
|||
|
|
|
|||
|
|
return unique
|
|||
|
|
|