599 lines
22 KiB
Python
599 lines
22 KiB
Python
"""
|
||
代码分析工具
|
||
使用 LLM 深度分析代码安全问题
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
from typing import Optional, List, Dict, Any
|
||
from pydantic import BaseModel, Field
|
||
|
||
from .base import AgentTool, ToolResult
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class CodeAnalysisInput(BaseModel):
|
||
"""代码分析输入"""
|
||
code: str = Field(description="要分析的代码内容")
|
||
file_path: str = Field(default="unknown", description="文件路径")
|
||
language: str = Field(default="python", description="编程语言")
|
||
focus: Optional[str] = Field(
|
||
default=None,
|
||
description="重点关注的漏洞类型，如 sql_injection, xss, command_injection"
|
||
)
|
||
context: Optional[str] = Field(
|
||
default=None,
|
||
description="额外的上下文信息，如相关的其他代码片段"
|
||
)
|
||
|
||
|
||
class CodeAnalysisTool(AgentTool):
|
||
"""
|
||
代码分析工具
|
||
使用 LLM 对代码进行深度安全分析
|
||
"""
|
||
|
||
def __init__(self, llm_service):
|
||
"""
|
||
初始化代码分析工具
|
||
|
||
Args:
|
||
llm_service: LLM 服务实例
|
||
"""
|
||
super().__init__()
|
||
self.llm_service = llm_service
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "code_analysis"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return """深度分析代码安全问题。
|
||
使用 LLM 对代码进行全面的安全审计，识别潜在漏洞。
|
||
|
||
使用场景:
|
||
- 对疑似有问题的代码进行深入分析
|
||
- 分析复杂的业务逻辑漏洞
|
||
- 追踪数据流和污点传播
|
||
- 生成详细的漏洞报告和修复建议
|
||
|
||
输入:
|
||
- code: 要分析的代码
|
||
- file_path: 文件路径
|
||
- language: 编程语言
|
||
- focus: 可选，重点关注的漏洞类型
|
||
- context: 可选，额外的上下文代码
|
||
|
||
这个工具会消耗较多的 Token，建议在确认有疑似问题后使用。"""
|
||
|
||
@property
|
||
def args_schema(self):
|
||
return CodeAnalysisInput
|
||
|
||
async def _execute(
|
||
self,
|
||
code: str,
|
||
file_path: str = "unknown",
|
||
language: str = "python",
|
||
focus: Optional[str] = None,
|
||
context: Optional[str] = None,
|
||
**kwargs
|
||
) -> ToolResult:
|
||
"""执行代码分析"""
|
||
import asyncio
|
||
|
||
try:
|
||
# 限制代码长度，避免超时
|
||
max_code_length = 50000 # 约 50KB
|
||
if len(code) > max_code_length:
|
||
code = code[:max_code_length] + "\n\n... (代码已截断，仅分析前 50000 字符)"
|
||
|
||
# 添加超时保护（5分钟）
|
||
try:
|
||
analysis = await asyncio.wait_for(
|
||
self.llm_service.analyze_code(code, language),
|
||
timeout=300.0 # 5分钟超时
|
||
)
|
||
except asyncio.TimeoutError:
|
||
return ToolResult(
|
||
success=False,
|
||
error="代码分析超时（超过5分钟）。代码可能过长或过于复杂，请尝试分析较小的代码片段。",
|
||
)
|
||
|
||
issues = analysis.get("issues", [])
|
||
|
||
if not issues:
|
||
return ToolResult(
|
||
success=True,
|
||
data="代码分析完成，未发现明显的安全问题。\n\n"
|
||
f"质量评分: {analysis.get('quality_score', 'N/A')}\n"
|
||
f"文件: {file_path}",
|
||
metadata={
|
||
"file_path": file_path,
|
||
"issues_count": 0,
|
||
"quality_score": analysis.get("quality_score"),
|
||
}
|
||
)
|
||
|
||
# 格式化输出
|
||
output_parts = [f"🔍 代码分析结果 - {file_path}\n"]
|
||
output_parts.append(f"发现 {len(issues)} 个问题:\n")
|
||
|
||
for i, issue in enumerate(issues):
|
||
severity_icon = {
|
||
"critical": "🔴",
|
||
"high": "🟠",
|
||
"medium": "🟡",
|
||
"low": "🟢"
|
||
}.get(issue.get("severity", ""), "⚪")
|
||
|
||
output_parts.append(f"\n{severity_icon} 问题 {i+1}: {issue.get('title', 'Unknown')}")
|
||
output_parts.append(f" 类型: {issue.get('type', 'unknown')}")
|
||
output_parts.append(f" 严重程度: {issue.get('severity', 'unknown')}")
|
||
output_parts.append(f" 行号: {issue.get('line', 'N/A')}")
|
||
output_parts.append(f" 描述: {issue.get('description', '')}")
|
||
|
||
if issue.get("code_snippet"):
|
||
output_parts.append(f" 代码片段:\n ```\n {issue.get('code_snippet')}\n ```")
|
||
|
||
if issue.get("suggestion"):
|
||
output_parts.append(f" 修复建议: {issue.get('suggestion')}")
|
||
|
||
if issue.get("ai_explanation"):
|
||
output_parts.append(f" AI解释: {issue.get('ai_explanation')}")
|
||
|
||
output_parts.append(f"\n质量评分: {analysis.get('quality_score', 'N/A')}/100")
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data="\n".join(output_parts),
|
||
metadata={
|
||
"file_path": file_path,
|
||
"issues_count": len(issues),
|
||
"quality_score": analysis.get("quality_score"),
|
||
"issues": issues,
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
logger.error(f"代码分析失败: {e}")
|
||
logger.error(f"LLM Provider: {self.llm_service.config.provider.value if self.llm_service.config else 'N/A'}")
|
||
logger.error(f"LLM Model: {self.llm_service.config.model if self.llm_service.config else 'N/A'}")
|
||
logger.error(f"API Key 前缀: {self.llm_service.config.api_key[:10] + '...' if self.llm_service.config and self.llm_service.config.api_key else 'N/A'}")
|
||
logger.error(traceback.format_exc())
|
||
return ToolResult(
|
||
success=False,
|
||
error=f"代码分析失败: {str(e)}",
|
||
)
|
||
|
||
|
||
class DataFlowAnalysisInput(BaseModel):
|
||
"""数据流分析输入"""
|
||
source_code: str = Field(description="包含数据源的代码")
|
||
sink_code: Optional[str] = Field(default=None, description="包含数据汇的代码（如危险函数）")
|
||
variable_name: str = Field(description="要追踪的变量名")
|
||
file_path: str = Field(default="unknown", description="文件路径")
|
||
|
||
|
||
class DataFlowAnalysisTool(AgentTool):
|
||
"""
|
||
数据流分析工具
|
||
追踪变量从源到汇的数据流
|
||
"""
|
||
|
||
def __init__(self, llm_service):
|
||
super().__init__()
|
||
self.llm_service = llm_service
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "dataflow_analysis"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return """分析代码中的数据流，追踪变量从源（如用户输入）到汇（如危险函数）的路径。
|
||
|
||
使用场景:
|
||
- 追踪用户输入如何流向危险函数
|
||
- 分析变量是否经过净化处理
|
||
- 识别污点传播路径
|
||
|
||
输入:
|
||
- source_code: 包含数据源的代码
|
||
- sink_code: 包含数据汇的代码（可选）
|
||
- variable_name: 要追踪的变量名
|
||
- file_path: 文件路径"""
|
||
|
||
@property
|
||
def args_schema(self):
|
||
return DataFlowAnalysisInput
|
||
|
||
async def _execute(
|
||
self,
|
||
source_code: str,
|
||
variable_name: str,
|
||
sink_code: Optional[str] = None,
|
||
file_path: str = "unknown",
|
||
**kwargs
|
||
) -> ToolResult:
|
||
"""执行数据流分析 - 增强版，带超时保护和回退逻辑"""
|
||
import asyncio
|
||
import re
|
||
|
||
# 🔥 首先尝试基于规则的快速分析（不依赖 LLM）
|
||
quick_analysis = self._quick_pattern_analysis(source_code, variable_name, sink_code)
|
||
|
||
try:
|
||
# 构建分析 prompt
|
||
analysis_prompt = f"""分析以下代码中变量 '{variable_name}' 的数据流。
|
||
|
||
源代码:
|
||
```
|
||
{source_code}
|
||
```
|
||
"""
|
||
if sink_code:
|
||
analysis_prompt += f"""
|
||
汇代码（可能的危险函数）:
|
||
```
|
||
{sink_code}
|
||
```
|
||
"""
|
||
|
||
analysis_prompt += f"""
|
||
请分析:
|
||
1. 变量 '{variable_name}' 的来源是什么？（用户输入、配置、数据库等）
|
||
2. 变量在传递过程中是否经过了净化/验证？
|
||
3. 变量最终流向了哪些危险函数？
|
||
4. 是否存在安全风险？
|
||
|
||
请返回 JSON 格式的分析结果，包含:
|
||
- source_type: 数据源类型
|
||
- sanitized: 是否经过净化
|
||
- sanitization_methods: 使用的净化方法
|
||
- dangerous_sinks: 流向的危险函数列表
|
||
- risk_level: 风险等级 (high/medium/low/none)
|
||
- explanation: 详细解释
|
||
- recommendation: 建议
|
||
"""
|
||
|
||
# 🔥 添加超时保护（2分钟）
|
||
try:
|
||
result = await asyncio.wait_for(
|
||
self.llm_service.analyze_code_with_custom_prompt(
|
||
code=source_code,
|
||
language="text",
|
||
custom_prompt=analysis_prompt,
|
||
),
|
||
timeout=120.0 # 2分钟超时
|
||
)
|
||
except asyncio.TimeoutError:
|
||
logger.warning(f"数据流分析 LLM 调用超时，使用快速分析结果")
|
||
return self._format_quick_analysis_result(quick_analysis, variable_name, file_path, "LLM调用超时，使用规则分析")
|
||
|
||
# 🔥 检查结果是否有效
|
||
if not result or (isinstance(result, dict) and not result.get("source_type") and not result.get("risk_level")):
|
||
logger.warning(f"数据流分析 LLM 返回无效结果，使用快速分析结果")
|
||
return self._format_quick_analysis_result(quick_analysis, variable_name, file_path, "LLM返回无效，使用规则分析")
|
||
|
||
# 格式化输出
|
||
output_parts = [f"📊 数据流分析结果 - 变量: {variable_name}\n"]
|
||
|
||
if isinstance(result, dict):
|
||
if result.get("source_type"):
|
||
output_parts.append(f"数据源: {result.get('source_type')}")
|
||
if result.get("sanitized") is not None:
|
||
sanitized = "✅ 是" if result.get("sanitized") else "❌ 否"
|
||
output_parts.append(f"是否净化: {sanitized}")
|
||
if result.get("sanitization_methods"):
|
||
methods = result.get('sanitization_methods', [])
|
||
if isinstance(methods, list):
|
||
output_parts.append(f"净化方法: {', '.join(methods)}")
|
||
else:
|
||
output_parts.append(f"净化方法: {methods}")
|
||
if result.get("dangerous_sinks"):
|
||
sinks = result.get('dangerous_sinks', [])
|
||
if isinstance(sinks, list):
|
||
output_parts.append(f"危险函数: {', '.join(sinks)}")
|
||
else:
|
||
output_parts.append(f"危险函数: {sinks}")
|
||
if result.get("risk_level"):
|
||
risk_icons = {"high": "🔴", "medium": "🟠", "low": "🟡", "none": "🟢"}
|
||
icon = risk_icons.get(result.get("risk_level", ""), "⚪")
|
||
output_parts.append(f"风险等级: {icon} {result.get('risk_level', '').upper()}")
|
||
if result.get("explanation"):
|
||
output_parts.append(f"\n分析: {result.get('explanation')}")
|
||
if result.get("recommendation"):
|
||
output_parts.append(f"\n建议: {result.get('recommendation')}")
|
||
else:
|
||
output_parts.append(str(result))
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data="\n".join(output_parts),
|
||
metadata={
|
||
"variable": variable_name,
|
||
"file_path": file_path,
|
||
"analysis": result,
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据流分析失败: {e}")
|
||
# 🔥 回退到快速分析
|
||
return self._format_quick_analysis_result(
|
||
quick_analysis,
|
||
variable_name,
|
||
file_path,
|
||
f"LLM调用失败({str(e)[:50]}...)，使用规则分析"
|
||
)
|
||
|
||
def _quick_pattern_analysis(
|
||
self,
|
||
source_code: str,
|
||
variable_name: str,
|
||
sink_code: Optional[str] = None
|
||
) -> Dict[str, Any]:
|
||
"""基于规则的快速数据流分析（不依赖 LLM）"""
|
||
import re
|
||
|
||
result = {
|
||
"source_type": "unknown",
|
||
"sanitized": False,
|
||
"sanitization_methods": [],
|
||
"dangerous_sinks": [],
|
||
"risk_level": "low",
|
||
}
|
||
|
||
code_to_analyze = source_code + (sink_code or "")
|
||
|
||
# 检测数据源类型
|
||
source_patterns = {
|
||
"user_input_get": r'\$_GET\[',
|
||
"user_input_post": r'\$_POST\[',
|
||
"user_input_request": r'\$_REQUEST\[',
|
||
"user_input_cookie": r'\$_COOKIE\[',
|
||
"request_param": r'request\.(GET|POST|args|form|data)',
|
||
"input_func": r'\binput\s*\(',
|
||
}
|
||
|
||
for source_name, pattern in source_patterns.items():
|
||
if re.search(pattern, source_code, re.IGNORECASE):
|
||
result["source_type"] = source_name
|
||
break
|
||
|
||
# 检测净化方法
|
||
sanitize_patterns = [
|
||
(r'htmlspecialchars\s*\(', "htmlspecialchars"),
|
||
(r'mysqli_real_escape_string\s*\(', "mysqli_escape"),
|
||
(r'addslashes\s*\(', "addslashes"),
|
||
(r'strip_tags\s*\(', "strip_tags"),
|
||
(r'filter_var\s*\(', "filter_var"),
|
||
(r'escape\s*\(', "escape"),
|
||
(r'sanitize', "sanitize"),
|
||
(r'validate', "validate"),
|
||
]
|
||
|
||
for pattern, name in sanitize_patterns:
|
||
if re.search(pattern, code_to_analyze, re.IGNORECASE):
|
||
result["sanitized"] = True
|
||
result["sanitization_methods"].append(name)
|
||
|
||
# 检测危险 sink
|
||
sink_patterns = [
|
||
(r'mysql_query\s*\(', "mysql_query"),
|
||
(r'mysqli_query\s*\(', "mysqli_query"),
|
||
(r'execute\s*\(', "execute"),
|
||
(r'shell_exec\s*\(', "shell_exec"),
|
||
(r'system\s*\(', "system"),
|
||
(r'exec\s*\(', "exec"),
|
||
(r'eval\s*\(', "eval"),
|
||
(r'include\s*\(', "include"),
|
||
(r'require\s*\(', "require"),
|
||
(r'file_get_contents\s*\(', "file_get_contents"),
|
||
(r'echo\s+', "echo"),
|
||
(r'print\s*\(', "print"),
|
||
]
|
||
|
||
for pattern, name in sink_patterns:
|
||
if re.search(pattern, code_to_analyze, re.IGNORECASE):
|
||
result["dangerous_sinks"].append(name)
|
||
|
||
# 计算风险等级
|
||
if result["source_type"].startswith("user_input") and result["dangerous_sinks"]:
|
||
if not result["sanitized"]:
|
||
result["risk_level"] = "high"
|
||
else:
|
||
result["risk_level"] = "medium"
|
||
elif result["dangerous_sinks"]:
|
||
result["risk_level"] = "medium"
|
||
|
||
return result
|
||
|
||
def _format_quick_analysis_result(
|
||
self,
|
||
analysis: Dict[str, Any],
|
||
variable_name: str,
|
||
file_path: str,
|
||
note: str
|
||
) -> ToolResult:
|
||
"""格式化快速分析结果"""
|
||
output_parts = [f"📊 数据流分析结果 - 变量: {variable_name}"]
|
||
output_parts.append(f"⚠️ 注意: {note}\n")
|
||
|
||
output_parts.append(f"数据源: {analysis.get('source_type', 'unknown')}")
|
||
output_parts.append(f"是否净化: {'✅ 是' if analysis.get('sanitized') else '❌ 否'}")
|
||
|
||
if analysis.get("sanitization_methods"):
|
||
output_parts.append(f"净化方法: {', '.join(analysis['sanitization_methods'])}")
|
||
|
||
if analysis.get("dangerous_sinks"):
|
||
output_parts.append(f"危险函数: {', '.join(analysis['dangerous_sinks'])}")
|
||
|
||
risk_icons = {"high": "🔴", "medium": "🟠", "low": "🟡", "none": "🟢"}
|
||
risk = analysis.get("risk_level", "low")
|
||
output_parts.append(f"风险等级: {risk_icons.get(risk, '⚪')} {risk.upper()}")
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data="\n".join(output_parts),
|
||
metadata={
|
||
"variable": variable_name,
|
||
"file_path": file_path,
|
||
"analysis": analysis,
|
||
"fallback_used": True,
|
||
}
|
||
)
|
||
|
||
|
||
class VulnerabilityValidationInput(BaseModel):
|
||
"""漏洞验证输入"""
|
||
code: str = Field(description="可能存在漏洞的代码")
|
||
vulnerability_type: str = Field(description="漏洞类型")
|
||
file_path: str = Field(default="unknown", description="文件路径")
|
||
line_number: Optional[int] = Field(default=None, description="行号")
|
||
context: Optional[str] = Field(default=None, description="额外上下文")
|
||
|
||
|
||
class VulnerabilityValidationTool(AgentTool):
|
||
"""
|
||
漏洞验证工具
|
||
验证疑似漏洞是否真实存在
|
||
"""
|
||
|
||
def __init__(self, llm_service):
|
||
super().__init__()
|
||
self.llm_service = llm_service
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "vulnerability_validation"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return """验证疑似漏洞是否真实存在。
|
||
对发现的潜在漏洞进行深入分析，判断是否为真正的安全问题。
|
||
|
||
输入:
|
||
- code: 包含疑似漏洞的代码
|
||
- vulnerability_type: 漏洞类型（如 sql_injection, xss 等）
|
||
- file_path: 文件路径
|
||
- line_number: 可选，行号
|
||
- context: 可选，额外的上下文代码
|
||
|
||
输出:
|
||
- 验证结果（确认/可能/误报）
|
||
- 详细分析
|
||
- 利用条件
|
||
- PoC 思路（如果确认存在漏洞）"""
|
||
|
||
@property
|
||
def args_schema(self):
|
||
return VulnerabilityValidationInput
|
||
|
||
async def _execute(
|
||
self,
|
||
code: str,
|
||
vulnerability_type: str,
|
||
file_path: str = "unknown",
|
||
line_number: Optional[int] = None,
|
||
context: Optional[str] = None,
|
||
**kwargs
|
||
) -> ToolResult:
|
||
"""执行漏洞验证"""
|
||
try:
|
||
validation_prompt = f"""你是一个专业的安全研究员，请验证以下代码中是否真的存在 {vulnerability_type} 漏洞。
|
||
|
||
代码:
|
||
```
|
||
{code}
|
||
```
|
||
|
||
{f'额外上下文:' + chr(10) + '```' + chr(10) + context + chr(10) + '```' if context else ''}
|
||
|
||
请分析:
|
||
1. 这段代码是否真的存在 {vulnerability_type} 漏洞？
|
||
2. 漏洞的利用条件是什么？
|
||
3. 攻击者如何利用这个漏洞？
|
||
4. 这是否可能是误报？为什么？
|
||
|
||
请返回 JSON 格式:
|
||
{{
|
||
"is_vulnerable": true/false/null (null表示无法确定),
|
||
"confidence": 0.0-1.0,
|
||
"verdict": "confirmed/likely/unlikely/false_positive",
|
||
"exploitation_conditions": ["条件1", "条件2"],
|
||
"attack_vector": "攻击向量描述",
|
||
"poc_idea": "PoC思路（如果存在漏洞）",
|
||
"false_positive_reason": "如果是误报，说明原因",
|
||
"detailed_analysis": "详细分析"
|
||
}}
|
||
"""
|
||
|
||
result = await self.llm_service.analyze_code_with_custom_prompt(
|
||
code=code,
|
||
language="text",
|
||
custom_prompt=validation_prompt,
|
||
)
|
||
|
||
# 格式化输出
|
||
output_parts = [f"🔎 漏洞验证结果 - {vulnerability_type}\n"]
|
||
output_parts.append(f"文件: {file_path}")
|
||
if line_number:
|
||
output_parts.append(f"行号: {line_number}")
|
||
output_parts.append("")
|
||
|
||
if isinstance(result, dict):
|
||
# 验证结果
|
||
verdict_icons = {
|
||
"confirmed": "🔴 确认存在漏洞",
|
||
"likely": "🟠 可能存在漏洞",
|
||
"unlikely": "🟡 可能是误报",
|
||
"false_positive": "🟢 误报",
|
||
}
|
||
verdict = result.get("verdict", "unknown")
|
||
output_parts.append(f"判定: {verdict_icons.get(verdict, verdict)}")
|
||
|
||
if result.get("confidence"):
|
||
output_parts.append(f"置信度: {result.get('confidence') * 100:.0f}%")
|
||
|
||
if result.get("exploitation_conditions"):
|
||
output_parts.append(f"\n利用条件:")
|
||
for cond in result.get("exploitation_conditions", []):
|
||
output_parts.append(f" - {cond}")
|
||
|
||
if result.get("attack_vector"):
|
||
output_parts.append(f"\n攻击向量: {result.get('attack_vector')}")
|
||
|
||
if result.get("poc_idea") and verdict in ["confirmed", "likely"]:
|
||
output_parts.append(f"\nPoC思路: {result.get('poc_idea')}")
|
||
|
||
if result.get("false_positive_reason") and verdict in ["unlikely", "false_positive"]:
|
||
output_parts.append(f"\n误报原因: {result.get('false_positive_reason')}")
|
||
|
||
if result.get("detailed_analysis"):
|
||
output_parts.append(f"\n详细分析:\n{result.get('detailed_analysis')}")
|
||
else:
|
||
output_parts.append(str(result))
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data="\n".join(output_parts),
|
||
metadata={
|
||
"vulnerability_type": vulnerability_type,
|
||
"file_path": file_path,
|
||
"line_number": line_number,
|
||
"validation": result,
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
return ToolResult(
|
||
success=False,
|
||
error=f"漏洞验证失败: {str(e)}",
|
||
)
|
||
|