267 lines
8.7 KiB
Python
267 lines
8.7 KiB
Python
|
"""
|
|||
|
|
自定义节点验证器
|
|||
|
|
检查节点代码的安全性和规范性
|
|||
|
|
"""
|
|||
|
|
import ast
|
|||
|
|
import re
|
|||
|
|
from typing import Dict, List, Optional, Tuple
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
|
|||
|
|
class NodeValidationError(Exception):
|
|||
|
|
"""节点验证错误"""
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
|
|||
|
|
class NodeValidator:
|
|||
|
|
"""节点代码验证器"""
|
|||
|
|
|
|||
|
|
# 危险的模块和函数（黑名单）
|
|||
|
|
DANGEROUS_MODULES = {
|
|||
|
|
'os', 'subprocess', 'sys', 'importlib', 'builtins',
|
|||
|
|
'__import__', 'eval', 'exec', 'compile', 'open',
|
|||
|
|
'file', 'input', 'raw_input'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 危险的函数调用
|
|||
|
|
DANGEROUS_CALLS = {
|
|||
|
|
'eval', 'exec', 'compile', '__import__',
|
|||
|
|
'system', 'popen', 'spawn', 'fork'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 必须的基类
|
|||
|
|
REQUIRED_BASE_CLASS = 'TraceNode'
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self.errors: List[str] = []
|
|||
|
|
self.warnings: List[str] = []
|
|||
|
|
|
|||
|
|
def validate_code(self, code: str, filename: str = '<string>') -> Dict:
|
|||
|
|
"""
|
|||
|
|
验证节点代码
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
code: Python代码字符串
|
|||
|
|
filename: 文件名（用于错误提示）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
{
|
|||
|
|
'valid': bool,
|
|||
|
|
'errors': List[str],
|
|||
|
|
'warnings': List[str],
|
|||
|
|
'node_classes': List[str], # 检测到的节点类
|
|||
|
|
'metadata': Dict # 提取的元数据
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
self.errors = []
|
|||
|
|
self.warnings = []
|
|||
|
|
node_classes = []
|
|||
|
|
|
|||
|
|
# 1. 基础语法检查
|
|||
|
|
try:
|
|||
|
|
tree = ast.parse(code, filename=filename)
|
|||
|
|
except SyntaxError as e:
|
|||
|
|
return {
|
|||
|
|
'valid': False,
|
|||
|
|
'errors': [f"语法错误 (行 {e.lineno}): {e.msg}"],
|
|||
|
|
'warnings': [],
|
|||
|
|
'node_classes': [],
|
|||
|
|
'metadata': {}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 2. AST安全检查
|
|||
|
|
self._check_security(tree)
|
|||
|
|
|
|||
|
|
# 3. 检查节点类定义
|
|||
|
|
node_classes = self._check_node_classes(tree)
|
|||
|
|
|
|||
|
|
# 4. 提取元数据
|
|||
|
|
metadata = self._extract_metadata(tree)
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'valid': len(self.errors) == 0,
|
|||
|
|
'errors': self.errors,
|
|||
|
|
'warnings': self.warnings,
|
|||
|
|
'node_classes': node_classes,
|
|||
|
|
'metadata': metadata
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def _check_security(self, tree: ast.AST):
|
|||
|
|
"""检查代码安全性"""
|
|||
|
|
for node in ast.walk(tree):
|
|||
|
|
# 检查危险的导入
|
|||
|
|
if isinstance(node, ast.Import):
|
|||
|
|
for alias in node.names:
|
|||
|
|
if alias.name in self.DANGEROUS_MODULES:
|
|||
|
|
self.errors.append(
|
|||
|
|
f"❌ 禁止导入危险模块: {alias.name} (行 {node.lineno})"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
elif isinstance(node, ast.ImportFrom):
|
|||
|
|
if node.module in self.DANGEROUS_MODULES:
|
|||
|
|
self.errors.append(
|
|||
|
|
f"❌ 禁止从危险模块导入: {node.module} (行 {node.lineno})"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 检查危险的函数调用
|
|||
|
|
elif isinstance(node, ast.Call):
|
|||
|
|
if isinstance(node.func, ast.Name):
|
|||
|
|
if node.func.id in self.DANGEROUS_CALLS:
|
|||
|
|
self.errors.append(
|
|||
|
|
f"❌ 禁止调用危险函数: {node.func.id} (行 {node.lineno})"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 检查文件操作
|
|||
|
|
elif isinstance(node, ast.Call):
|
|||
|
|
if isinstance(node.func, ast.Name) and node.func.id == 'open':
|
|||
|
|
self.warnings.append(
|
|||
|
|
f"⚠️ 检测到文件操作 (行 {node.lineno})，请确保路径安全"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def _check_node_classes(self, tree: ast.AST) -> List[str]:
|
|||
|
|
"""检查节点类定义"""
|
|||
|
|
node_classes = []
|
|||
|
|
|
|||
|
|
for node in ast.walk(tree):
|
|||
|
|
if isinstance(node, ast.ClassDef):
|
|||
|
|
# 检查是否继承TraceNode
|
|||
|
|
has_trace_base = False
|
|||
|
|
for base in node.bases:
|
|||
|
|
if isinstance(base, ast.Name) and base.id == self.REQUIRED_BASE_CLASS:
|
|||
|
|
has_trace_base = True
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
if has_trace_base:
|
|||
|
|
node_classes.append(node.name)
|
|||
|
|
|
|||
|
|
# 检查必须的方法
|
|||
|
|
methods = {n.name for n in node.body if isinstance(n, ast.FunctionDef)}
|
|||
|
|
|
|||
|
|
if 'execute' not in methods:
|
|||
|
|
self.errors.append(
|
|||
|
|
f"❌ 节点类 {node.name} 缺少 execute() 方法 (行 {node.lineno})"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if 'get_metadata' not in methods:
|
|||
|
|
self.warnings.append(
|
|||
|
|
f"⚠️ 节点类 {node.name} 建议实现 get_metadata() 静态方法"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not node_classes:
|
|||
|
|
self.errors.append(
|
|||
|
|
f"❌ 未找到继承自 {self.REQUIRED_BASE_CLASS} 的节点类"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return node_classes
|
|||
|
|
|
|||
|
|
def _extract_metadata(self, tree: ast.AST) -> Dict:
|
|||
|
|
"""提取节点元数据"""
|
|||
|
|
metadata = {}
|
|||
|
|
|
|||
|
|
for node in ast.walk(tree):
|
|||
|
|
if isinstance(node, ast.ClassDef):
|
|||
|
|
for item in node.body:
|
|||
|
|
if isinstance(item, ast.FunctionDef) and item.name == 'get_metadata':
|
|||
|
|
# 尝试提取返回的字典
|
|||
|
|
for stmt in item.body:
|
|||
|
|
if isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.Dict):
|
|||
|
|
try:
|
|||
|
|
# 简单提取（只处理字面量）
|
|||
|
|
metadata = ast.literal_eval(stmt.value)
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
return metadata
|
|||
|
|
|
|||
|
|
def validate_file(self, filepath: Path) -> Dict:
|
|||
|
|
"""验证节点文件"""
|
|||
|
|
try:
|
|||
|
|
code = filepath.read_text(encoding='utf-8')
|
|||
|
|
return self.validate_code(code, str(filepath))
|
|||
|
|
except Exception as e:
|
|||
|
|
return {
|
|||
|
|
'valid': False,
|
|||
|
|
'errors': [f"读取文件失败: {str(e)}"],
|
|||
|
|
'warnings': [],
|
|||
|
|
'node_classes': [],
|
|||
|
|
'metadata': {}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def is_safe_filename(filename: str) -> Tuple[bool, Optional[str]]:
|
|||
|
|
"""
|
|||
|
|
检查文件名是否安全
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
(is_safe, error_message)
|
|||
|
|
"""
|
|||
|
|
# 必须是.py文件
|
|||
|
|
if not filename.endswith('.py'):
|
|||
|
|
return False, "文件必须是 .py 扩展名"
|
|||
|
|
|
|||
|
|
# 不能以__开头（Python内部文件）
|
|||
|
|
name = filename[:-3]
|
|||
|
|
if name.startswith('__'):
|
|||
|
|
return False, "文件名不能以 __ 开头"
|
|||
|
|
|
|||
|
|
# 不能包含特殊字符
|
|||
|
|
if not re.match(r'^[a-zA-Z0-9_]+\.py$', filename):
|
|||
|
|
return False, "文件名只能包含字母、数字和下划线"
|
|||
|
|
|
|||
|
|
# 保留的文件名
|
|||
|
|
reserved = {'__init__', 'test', 'setup', 'config'}
|
|||
|
|
if name.lower() in reserved:
|
|||
|
|
return False, f"文件名 '{name}' 是保留名称"
|
|||
|
|
|
|||
|
|
return True, None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class NodeSandbox:
|
|||
|
|
"""节点执行沙箱（简化版）"""
|
|||
|
|
|
|||
|
|
MAX_EXECUTION_TIME = 30 # 秒
|
|||
|
|
MAX_MEMORY_MB = 512
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def execute_with_limits(node_instance, inputs: Dict, timeout: int = None):
|
|||
|
|
"""
|
|||
|
|
在资源限制下执行节点
|
|||
|
|
|
|||
|
|
TODO: 实现真正的沙箱（使用 RestrictedPython 或容器）
|
|||
|
|
"""
|
|||
|
|
import time
|
|||
|
|
import threading
|
|||
|
|
|
|||
|
|
timeout = timeout or NodeSandbox.MAX_EXECUTION_TIME
|
|||
|
|
result = {'output': None, 'error': None, 'timeout': False}
|
|||
|
|
|
|||
|
|
def target():
|
|||
|
|
try:
|
|||
|
|
result['output'] = node_instance.execute(inputs)
|
|||
|
|
except Exception as e:
|
|||
|
|
result['error'] = str(e)
|
|||
|
|
|
|||
|
|
thread = threading.Thread(target=target)
|
|||
|
|
thread.daemon = True
|
|||
|
|
thread.start()
|
|||
|
|
thread.join(timeout)
|
|||
|
|
|
|||
|
|
if thread.is_alive():
|
|||
|
|
result['timeout'] = True
|
|||
|
|
result['error'] = f"节点执行超时 ({timeout}秒)"
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 便捷函数
|
|||
|
|
def validate_node_code(code: str) -> Dict:
|
|||
|
|
"""验证节点代码（快捷方式）"""
|
|||
|
|
validator = NodeValidator()
|
|||
|
|
return validator.validate_code(code)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def validate_node_file(filepath: str) -> Dict:
|
|||
|
|
"""验证节点文件（快捷方式）"""
|
|||
|
|
validator = NodeValidator()
|
|||
|
|
return validator.validate_file(Path(filepath))
|