514 lines
17 KiB
Python
514 lines
17 KiB
Python
"""
|
||
通用代码执行工具 - LLM 驱动的漏洞验证
|
||
|
||
核心理念：
|
||
- LLM 是验证的大脑，工具只提供执行能力
|
||
- 不硬编码 payload、检测规则
|
||
- LLM 自己决定测试策略、编写测试代码、分析结果
|
||
|
||
使用场景：
|
||
- LLM 编写 Fuzzing Harness 进行局部测试
|
||
- LLM 构造 PoC 验证漏洞
|
||
- LLM 编写 mock 代码隔离测试函数
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import tempfile
|
||
from typing import Optional, Dict, Any
|
||
from pydantic import BaseModel, Field
|
||
|
||
from .base import AgentTool, ToolResult
|
||
from .sandbox_tool import SandboxManager, SandboxConfig
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class RunCodeInput(BaseModel):
|
||
"""代码执行输入"""
|
||
code: str = Field(..., description="要执行的代码")
|
||
language: str = Field(default="python", description="编程语言: python, php, javascript, ruby, go, java, bash")
|
||
timeout: int = Field(default=60, description="超时时间（秒），复杂测试可设置更长")
|
||
description: str = Field(default="", description="简短描述这段代码的目的（用于日志）")
|
||
|
||
|
||
class RunCodeTool(AgentTool):
|
||
"""
|
||
通用代码执行工具
|
||
|
||
让 LLM 自由编写测试代码，在沙箱中执行。
|
||
|
||
LLM 可以：
|
||
- 编写 Fuzzing Harness 隔离测试单个函数
|
||
- 构造 mock 对象模拟依赖
|
||
- 设计各种 payload 进行测试
|
||
- 分析执行结果判断漏洞
|
||
|
||
工具不做任何假设，完全由 LLM 控制测试逻辑。
|
||
"""
|
||
|
||
def __init__(self, sandbox_manager: Optional[SandboxManager] = None, project_root: str = "."):
|
||
super().__init__()
|
||
# 使用更宽松的沙箱配置
|
||
config = SandboxConfig(
|
||
timeout=120,
|
||
memory_limit="1g", # 更大内存
|
||
)
|
||
self.sandbox_manager = sandbox_manager or SandboxManager(config)
|
||
self.project_root = project_root
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "run_code"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return """🔥 通用代码执行工具 - 在沙箱中运行你编写的测试代码
|
||
|
||
这是你进行漏洞验证的核心工具。你可以：
|
||
1. 编写 Fuzzing Harness 隔离测试单个函数
|
||
2. 构造 mock 对象模拟数据库、HTTP 请求等依赖
|
||
3. 设计各种 payload 进行漏洞测试
|
||
4. 编写完整的 PoC 验证脚本
|
||
|
||
输入：
|
||
- code: 你编写的测试代码（完整可执行）
|
||
- language: python, php, javascript, ruby, go, java, bash
|
||
- timeout: 超时秒数（默认60，复杂测试可设更长）
|
||
- description: 简短描述代码目的
|
||
|
||
支持的语言和执行方式：
|
||
- python: python3 -c 'code'
|
||
- php: php -r 'code' (注意：不需要 <?php 标签)
|
||
- javascript: node -e 'code'
|
||
- ruby: ruby -e 'code'
|
||
- go: go run (需写完整 package main)
|
||
- java: javac + java (需写完整 class)
|
||
- bash: bash -c 'code'
|
||
|
||
示例 - 命令注入 Fuzzing Harness:
|
||
```python
|
||
# 提取目标函数并构造测试
|
||
import os
|
||
|
||
# Mock os.system 来检测是否被调用
|
||
executed_commands = []
|
||
original_system = os.system
|
||
def mock_system(cmd):
|
||
print(f"[DETECTED] os.system called: {cmd}")
|
||
executed_commands.append(cmd)
|
||
return 0
|
||
os.system = mock_system
|
||
|
||
# 目标函数（从项目代码复制）
|
||
def vulnerable_function(user_input):
|
||
os.system(f"echo {user_input}")
|
||
|
||
# Fuzzing 测试
|
||
payloads = ["; id", "| whoami", "$(cat /etc/passwd)", "`id`"]
|
||
for payload in payloads:
|
||
print(f"\\nTesting payload: {payload}")
|
||
executed_commands.clear()
|
||
try:
|
||
vulnerable_function(payload)
|
||
if executed_commands:
|
||
print(f"[VULN] Command injection detected!")
|
||
except Exception as e:
|
||
print(f"Error: {e}")
|
||
```
|
||
|
||
⚠️ 重要提示：
|
||
- 代码在 Docker 沙箱中执行，与真实环境隔离
|
||
- 你需要自己 mock 依赖（数据库、HTTP、文件系统等）
|
||
- 你需要自己设计 payload 和检测逻辑
|
||
- 你需要自己分析输出判断漏洞是否存在"""
|
||
|
||
@property
|
||
def args_schema(self):
|
||
return RunCodeInput
|
||
|
||
async def _execute(
|
||
self,
|
||
code: str,
|
||
language: str = "python",
|
||
timeout: int = 60,
|
||
description: str = "",
|
||
**kwargs
|
||
) -> ToolResult:
|
||
"""执行用户编写的代码"""
|
||
|
||
# 初始化沙箱
|
||
try:
|
||
await self.sandbox_manager.initialize()
|
||
except Exception as e:
|
||
logger.warning(f"Sandbox init failed: {e}")
|
||
|
||
if not self.sandbox_manager.is_available:
|
||
return ToolResult(
|
||
success=False,
|
||
error="沙箱环境不可用 (Docker 未运行)",
|
||
data="请确保 Docker 已启动。如果无法使用沙箱，你可以通过静态分析代码来验证漏洞。"
|
||
)
|
||
|
||
# 构建执行命令
|
||
language = language.lower().strip()
|
||
command = self._build_command(code, language)
|
||
|
||
if command is None:
|
||
return ToolResult(
|
||
success=False,
|
||
error=f"不支持的语言: {language}",
|
||
data=f"支持的语言: python, php, javascript, ruby, go, java, bash"
|
||
)
|
||
|
||
# 在沙箱中执行
|
||
result = await self.sandbox_manager.execute_command(
|
||
command=command,
|
||
timeout=timeout,
|
||
)
|
||
|
||
# 格式化输出
|
||
output_parts = [f"🔬 代码执行结果"]
|
||
if description:
|
||
output_parts.append(f"目的: {description}")
|
||
output_parts.append(f"语言: {language}")
|
||
output_parts.append(f"退出码: {result['exit_code']}")
|
||
|
||
if result.get("stdout"):
|
||
stdout = result["stdout"]
|
||
if len(stdout) > 5000:
|
||
stdout = stdout[:5000] + f"\n... (截断，共 {len(result['stdout'])} 字符)"
|
||
output_parts.append(f"\n输出:\n```\n{stdout}\n```")
|
||
|
||
if result.get("stderr"):
|
||
stderr = result["stderr"]
|
||
if len(stderr) > 2000:
|
||
stderr = stderr[:2000] + "\n... (截断)"
|
||
output_parts.append(f"\n错误输出:\n```\n{stderr}\n```")
|
||
|
||
if result.get("error"):
|
||
output_parts.append(f"\n执行错误: {result['error']}")
|
||
|
||
# 提示 LLM 分析结果
|
||
output_parts.append("\n---")
|
||
output_parts.append("请根据上述输出分析漏洞是否存在。")
|
||
|
||
return ToolResult(
|
||
success=result.get("success", False),
|
||
data="\n".join(output_parts),
|
||
error=result.get("error"),
|
||
metadata={
|
||
"language": language,
|
||
"exit_code": result.get("exit_code", -1),
|
||
"stdout_length": len(result.get("stdout", "")),
|
||
"stderr_length": len(result.get("stderr", "")),
|
||
}
|
||
)
|
||
|
||
def _build_command(self, code: str, language: str) -> Optional[str]:
|
||
"""根据语言构建执行命令"""
|
||
|
||
# 转义单引号的通用方法
|
||
def escape_for_shell(s: str) -> str:
|
||
return s.replace("'", "'\"'\"'")
|
||
|
||
if language == "python":
|
||
escaped = escape_for_shell(code)
|
||
return f"python3 -c '{escaped}'"
|
||
|
||
elif language == "php":
|
||
# PHP: php -r 不需要 <?php 标签
|
||
clean_code = code.strip()
|
||
if clean_code.startswith("<?php"):
|
||
clean_code = clean_code[5:].strip()
|
||
if clean_code.startswith("<?"):
|
||
clean_code = clean_code[2:].strip()
|
||
if clean_code.endswith("?>"):
|
||
clean_code = clean_code[:-2].strip()
|
||
escaped = escape_for_shell(clean_code)
|
||
return f"php -r '{escaped}'"
|
||
|
||
elif language in ["javascript", "js", "node"]:
|
||
escaped = escape_for_shell(code)
|
||
return f"node -e '{escaped}'"
|
||
|
||
elif language == "ruby":
|
||
escaped = escape_for_shell(code)
|
||
return f"ruby -e '{escaped}'"
|
||
|
||
elif language == "bash":
|
||
escaped = escape_for_shell(code)
|
||
return f"bash -c '{escaped}'"
|
||
|
||
elif language == "go":
|
||
# Go 需要完整的 package main
|
||
escaped = escape_for_shell(code).replace("\\", "\\\\")
|
||
return f"echo '{escaped}' > /tmp/main.go && go run /tmp/main.go"
|
||
|
||
elif language == "java":
|
||
# Java 需要完整的 class
|
||
escaped = escape_for_shell(code).replace("\\", "\\\\")
|
||
# 提取类名
|
||
import re
|
||
class_match = re.search(r'public\s+class\s+(\w+)', code)
|
||
class_name = class_match.group(1) if class_match else "Test"
|
||
return f"echo '{escaped}' > /tmp/{class_name}.java && javac /tmp/{class_name}.java && java -cp /tmp {class_name}"
|
||
|
||
return None
|
||
|
||
|
||
class ExtractFunctionInput(BaseModel):
|
||
"""函数提取输入"""
|
||
file_path: str = Field(..., description="源文件路径")
|
||
function_name: str = Field(..., description="要提取的函数名")
|
||
include_imports: bool = Field(default=True, description="是否包含 import 语句")
|
||
|
||
|
||
class ExtractFunctionTool(AgentTool):
|
||
"""
|
||
函数提取工具
|
||
|
||
从源文件中提取指定函数及其依赖，用于构建 Fuzzing Harness
|
||
"""
|
||
|
||
def __init__(self, project_root: str = "."):
|
||
super().__init__()
|
||
self.project_root = project_root
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "extract_function"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return """从源文件中提取指定函数的代码
|
||
|
||
用于构建 Fuzzing Harness 时获取目标函数代码。
|
||
|
||
输入：
|
||
- file_path: 源文件路径
|
||
- function_name: 要提取的函数名
|
||
- include_imports: 是否包含文件开头的 import 语句（默认 true）
|
||
|
||
返回：
|
||
- 函数代码
|
||
- 相关的 import 语句
|
||
- 函数参数列表
|
||
|
||
示例：
|
||
{"file_path": "app/api.py", "function_name": "process_command"}"""
|
||
|
||
@property
|
||
def args_schema(self):
|
||
return ExtractFunctionInput
|
||
|
||
async def _execute(
|
||
self,
|
||
file_path: str,
|
||
function_name: str,
|
||
include_imports: bool = True,
|
||
**kwargs
|
||
) -> ToolResult:
|
||
"""提取函数代码"""
|
||
import ast
|
||
import re
|
||
|
||
full_path = os.path.join(self.project_root, file_path)
|
||
if not os.path.exists(full_path):
|
||
return ToolResult(success=False, error=f"文件不存在: {file_path}")
|
||
|
||
with open(full_path, 'r', encoding='utf-8', errors='ignore') as f:
|
||
code = f.read()
|
||
|
||
# 检测语言
|
||
ext = os.path.splitext(file_path)[1].lower()
|
||
|
||
if ext == ".py":
|
||
result = self._extract_python(code, function_name, include_imports)
|
||
elif ext == ".php":
|
||
result = self._extract_php(code, function_name)
|
||
elif ext in [".js", ".ts"]:
|
||
result = self._extract_javascript(code, function_name)
|
||
else:
|
||
result = self._extract_generic(code, function_name)
|
||
|
||
if result["success"]:
|
||
output_parts = [f"📦 函数提取结果\n"]
|
||
output_parts.append(f"文件: {file_path}")
|
||
output_parts.append(f"函数: {function_name}")
|
||
|
||
if result.get("imports"):
|
||
output_parts.append(f"\n相关 imports:\n```\n{result['imports']}\n```")
|
||
|
||
if result.get("parameters"):
|
||
output_parts.append(f"\n参数: {', '.join(result['parameters'])}")
|
||
|
||
output_parts.append(f"\n函数代码:\n```\n{result['code']}\n```")
|
||
|
||
output_parts.append("\n---")
|
||
output_parts.append("你现在可以使用这段代码构建 Fuzzing Harness")
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data="\n".join(output_parts),
|
||
metadata=result
|
||
)
|
||
else:
|
||
return ToolResult(
|
||
success=False,
|
||
error=result.get("error", "提取失败"),
|
||
data=f"无法提取函数 '{function_name}'。你可以使用 read_file 工具直接读取文件，手动定位函数代码。"
|
||
)
|
||
|
||
def _extract_python(self, code: str, function_name: str, include_imports: bool) -> Dict:
|
||
"""提取 Python 函数"""
|
||
import ast
|
||
|
||
try:
|
||
tree = ast.parse(code)
|
||
except SyntaxError:
|
||
# 降级到正则提取
|
||
return self._extract_generic(code, function_name)
|
||
|
||
# 收集 imports
|
||
imports = []
|
||
if include_imports:
|
||
for node in ast.walk(tree):
|
||
if isinstance(node, ast.Import):
|
||
imports.append(ast.unparse(node))
|
||
elif isinstance(node, ast.ImportFrom):
|
||
imports.append(ast.unparse(node))
|
||
|
||
# 查找函数
|
||
for node in ast.walk(tree):
|
||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||
if node.name == function_name:
|
||
lines = code.split('\n')
|
||
func_code = '\n'.join(lines[node.lineno - 1:node.end_lineno])
|
||
params = [arg.arg for arg in node.args.args]
|
||
|
||
return {
|
||
"success": True,
|
||
"code": func_code,
|
||
"imports": '\n'.join(imports) if imports else None,
|
||
"parameters": params,
|
||
"line_start": node.lineno,
|
||
"line_end": node.end_lineno,
|
||
}
|
||
|
||
return {"success": False, "error": f"未找到函数 '{function_name}'"}
|
||
|
||
def _extract_php(self, code: str, function_name: str) -> Dict:
|
||
"""提取 PHP 函数"""
|
||
import re
|
||
|
||
pattern = rf'function\s+{re.escape(function_name)}\s*\([^)]*\)\s*\{{'
|
||
match = re.search(pattern, code)
|
||
|
||
if not match:
|
||
return {"success": False, "error": f"未找到函数 '{function_name}'"}
|
||
|
||
start_pos = match.start()
|
||
brace_count = 0
|
||
end_pos = match.end() - 1
|
||
|
||
for i, char in enumerate(code[match.end() - 1:], start=match.end() - 1):
|
||
if char == '{':
|
||
brace_count += 1
|
||
elif char == '}':
|
||
brace_count -= 1
|
||
if brace_count == 0:
|
||
end_pos = i + 1
|
||
break
|
||
|
||
func_code = code[start_pos:end_pos]
|
||
|
||
# 提取参数
|
||
param_match = re.search(r'function\s+\w+\s*\(([^)]*)\)', func_code)
|
||
params = []
|
||
if param_match:
|
||
params_str = param_match.group(1)
|
||
params = [p.strip().split('=')[0].strip().replace('$', '')
|
||
for p in params_str.split(',') if p.strip()]
|
||
|
||
return {
|
||
"success": True,
|
||
"code": func_code,
|
||
"parameters": params,
|
||
}
|
||
|
||
def _extract_javascript(self, code: str, function_name: str) -> Dict:
|
||
"""提取 JavaScript 函数"""
|
||
import re
|
||
|
||
patterns = [
|
||
rf'function\s+{re.escape(function_name)}\s*\([^)]*\)\s*\{{',
|
||
rf'(?:const|let|var)\s+{re.escape(function_name)}\s*=\s*function\s*\([^)]*\)\s*\{{',
|
||
rf'(?:const|let|var)\s+{re.escape(function_name)}\s*=\s*\([^)]*\)\s*=>\s*\{{',
|
||
rf'async\s+function\s+{re.escape(function_name)}\s*\([^)]*\)\s*\{{',
|
||
]
|
||
|
||
for pattern in patterns:
|
||
match = re.search(pattern, code)
|
||
if match:
|
||
start_pos = match.start()
|
||
brace_count = 0
|
||
end_pos = match.end() - 1
|
||
|
||
for i, char in enumerate(code[match.end() - 1:], start=match.end() - 1):
|
||
if char == '{':
|
||
brace_count += 1
|
||
elif char == '}':
|
||
brace_count -= 1
|
||
if brace_count == 0:
|
||
end_pos = i + 1
|
||
break
|
||
|
||
func_code = code[start_pos:end_pos]
|
||
|
||
return {
|
||
"success": True,
|
||
"code": func_code,
|
||
}
|
||
|
||
return {"success": False, "error": f"未找到函数 '{function_name}'"}
|
||
|
||
def _extract_generic(self, code: str, function_name: str) -> Dict:
|
||
"""通用函数提取（正则）"""
|
||
import re
|
||
|
||
# 尝试多种模式
|
||
patterns = [
|
||
rf'def\s+{re.escape(function_name)}\s*\([^)]*\)\s*:', # Python
|
||
rf'function\s+{re.escape(function_name)}\s*\([^)]*\)', # PHP/JS
|
||
rf'func\s+{re.escape(function_name)}\s*\([^)]*\)', # Go
|
||
]
|
||
|
||
for pattern in patterns:
|
||
match = re.search(pattern, code, re.MULTILINE)
|
||
if match:
|
||
start_line = code[:match.start()].count('\n')
|
||
lines = code.split('\n')
|
||
|
||
# 尝试找到函数结束
|
||
end_line = start_line + 1
|
||
indent = len(lines[start_line]) - len(lines[start_line].lstrip())
|
||
|
||
for i in range(start_line + 1, min(start_line + 100, len(lines))):
|
||
line = lines[i]
|
||
if line.strip() and not line.startswith(' ' * (indent + 1)):
|
||
if not line.strip().startswith('#'):
|
||
end_line = i
|
||
break
|
||
end_line = i + 1
|
||
|
||
func_code = '\n'.join(lines[start_line:end_line])
|
||
|
||
return {
|
||
"success": True,
|
||
"code": func_code,
|
||
}
|
||
|
||
return {"success": False, "error": f"未找到函数 '{function_name}'"}
|