314 lines
10 KiB
Python
314 lines
10 KiB
Python
"""
|
||
Python代码规则执行器
|
||
|
||
负责安全地执行规则中包含的Python代码，确保代码在隔离的环境中运行，
|
||
并对可执行的操作进行限制，以防恶意代码。
|
||
"""
|
||
|
||
import ast
|
||
import logging
|
||
import importlib
|
||
import inspect
|
||
import time
|
||
import threading
|
||
import os
|
||
from typing import Dict, Any, Optional, List, Callable, Tuple
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
|
||
from ..models.rule_models import PythonCodeRule
|
||
from ..models.rule_models import RuleCategory
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class CodeExecutionError(Exception):
|
||
"""执行Python代码时发生的错误"""
|
||
pass
|
||
|
||
class TimeoutError(CodeExecutionError):
|
||
"""代码执行超时错误"""
|
||
pass
|
||
|
||
class ImportError(CodeExecutionError):
|
||
"""非法导入模块错误"""
|
||
pass
|
||
|
||
class ValidationResult:
|
||
"""Python代码验证的结果"""
|
||
def __init__(self,
|
||
is_valid: bool,
|
||
message: Optional[str] = None,
|
||
details: Optional[Dict[str, Any]] = None,
|
||
exception: Optional[Exception] = None):
|
||
self.is_valid = is_valid
|
||
self.message = message
|
||
self.details = details or {}
|
||
self.exception = exception
|
||
|
||
def __bool__(self):
|
||
return self.is_valid
|
||
|
||
class PythonRuleExecutor:
|
||
"""
|
||
Python代码规则执行器
|
||
|
||
负责安全地执行规则中的Python代码，并返回验证结果。
|
||
"""
|
||
|
||
def __init__(self, rules_base_path: str = "./rules"):
|
||
self.logger = logging.getLogger(__name__)
|
||
self.rules_base_path = os.path.abspath(rules_base_path)
|
||
|
||
def _load_code_from_file(self, rule: PythonCodeRule) -> str:
|
||
"""
|
||
从文件加载代码
|
||
|
||
Args:
|
||
rule: Python代码规则对象，包含code_file属性
|
||
|
||
Returns:
|
||
加载的代码内容
|
||
|
||
Raises:
|
||
CodeExecutionError: 如果加载代码失败
|
||
"""
|
||
if not rule.code_file:
|
||
raise CodeExecutionError("规则未指定code_file属性")
|
||
|
||
# 构建代码文件的绝对路径
|
||
# 如果规则ID和版本都存在，则在python_code目录下查找
|
||
if rule.id and rule.version:
|
||
# 代码文件路径格式: rules/python_code/{rule_id}/{version}.py
|
||
file_path = os.path.join(self.rules_base_path, "python_code", rule.id, f"{rule.version}.py")
|
||
else:
|
||
# 否则使用传入的相对路径
|
||
file_path = os.path.join(self.rules_base_path, rule.code_file)
|
||
|
||
# 检查文件是否存在
|
||
if not os.path.isfile(file_path):
|
||
raise CodeExecutionError(f"代码文件不存在: {file_path}")
|
||
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
return f.read()
|
||
except Exception as e:
|
||
raise CodeExecutionError(f"读取代码文件失败: {e}")
|
||
|
||
def analyze_code(self, code: str) -> List[str]:
|
||
"""
|
||
分析代码，检查其中使用的导入模块
|
||
|
||
Args:
|
||
code: 要分析的Python代码
|
||
|
||
Returns:
|
||
代码中导入的模块列表
|
||
|
||
Raises:
|
||
SyntaxError: 如果代码存在语法错误
|
||
"""
|
||
try:
|
||
tree = ast.parse(code)
|
||
except SyntaxError as e:
|
||
raise CodeExecutionError(f"代码语法错误: {e}")
|
||
|
||
imports = []
|
||
for node in ast.walk(tree):
|
||
if isinstance(node, ast.Import):
|
||
for name in node.names:
|
||
imports.append(name.name)
|
||
elif isinstance(node, ast.ImportFrom):
|
||
imports.append(node.module)
|
||
|
||
return imports
|
||
|
||
def _execute_with_timeout(self,
|
||
func: Callable,
|
||
args: Tuple,
|
||
kwargs: Dict[str, Any],
|
||
timeout: int) -> Any:
|
||
"""
|
||
使用超时执行函数
|
||
|
||
Args:
|
||
func: 要执行的函数
|
||
args: 函数的位置参数
|
||
kwargs: 函数的关键字参数
|
||
timeout: 超时时间（秒）
|
||
|
||
Returns:
|
||
函数的返回值
|
||
|
||
Raises:
|
||
TimeoutError: 如果函数执行超过指定的超时时间
|
||
"""
|
||
result = [None]
|
||
exception = [None]
|
||
|
||
def target():
|
||
try:
|
||
result[0] = func(*args, **kwargs)
|
||
except Exception as e:
|
||
exception[0] = e
|
||
|
||
thread = threading.Thread(target=target)
|
||
thread.daemon = True
|
||
|
||
thread.start()
|
||
thread.join(timeout)
|
||
|
||
if thread.is_alive():
|
||
# 超时
|
||
raise TimeoutError(f"代码执行超时（超过{timeout}秒）")
|
||
|
||
if exception[0]:
|
||
raise exception[0]
|
||
|
||
return result[0]
|
||
|
||
def execute_rule(self, rule: PythonCodeRule, context: Dict[str, Any]) -> ValidationResult:
|
||
"""
|
||
执行Python代码规则
|
||
|
||
Args:
|
||
rule: Python代码规则定义
|
||
context: 验证上下文（包含规则执行所需的参数和数据）
|
||
|
||
Returns:
|
||
ValidationResult: 验证结果
|
||
"""
|
||
# 检查是否提供了所有必需的参数
|
||
if rule.expected_parameters:
|
||
missing_params = [p for p in rule.expected_parameters if p not in context]
|
||
if missing_params:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=f"缺少必需的参数: {', '.join(missing_params)}"
|
||
)
|
||
|
||
# 获取代码内容
|
||
try:
|
||
if rule.code:
|
||
code = rule.code
|
||
elif rule.code_file:
|
||
code = self._load_code_from_file(rule)
|
||
else:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message="规则既没有code也没有code_file属性"
|
||
)
|
||
except CodeExecutionError as e:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=str(e),
|
||
exception=e
|
||
)
|
||
|
||
# 分析代码中的导入
|
||
try:
|
||
imports = self.analyze_code(code)
|
||
except CodeExecutionError as e:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=str(e),
|
||
exception=e
|
||
)
|
||
|
||
# 检查导入是否被允许
|
||
if imports and not rule.allow_imports:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=f"规则不允许导入模块，但代码尝试导入: {', '.join(imports)}"
|
||
)
|
||
|
||
# 如果允许导入，检查是否所有导入都在允许列表中
|
||
if imports and rule.allow_imports and rule.allowed_modules:
|
||
unauthorized_imports = [imp for imp in imports if imp not in rule.allowed_modules]
|
||
if unauthorized_imports:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=f"代码尝试导入未授权的模块: {', '.join(unauthorized_imports)}"
|
||
)
|
||
|
||
# 准备执行环境
|
||
# 创建一个隔离的命名空间
|
||
namespace = {'__builtins__': __builtins__}
|
||
|
||
# 添加上下文变量
|
||
namespace.update(context)
|
||
|
||
# 如果允许导入，预先导入允许的模块
|
||
if rule.allow_imports and imports:
|
||
for module_name in imports:
|
||
if not rule.allowed_modules or module_name in rule.allowed_modules:
|
||
try:
|
||
module = importlib.import_module(module_name)
|
||
namespace[module_name] = module
|
||
except Exception as e:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=f"导入模块 '{module_name}' 失败: {e}",
|
||
exception=e
|
||
)
|
||
|
||
# 执行代码
|
||
try:
|
||
# 编译代码
|
||
compiled_code = compile(code, f"<rule:{rule.id}>", 'exec')
|
||
|
||
# 执行代码
|
||
self._execute_with_timeout(
|
||
exec,
|
||
(compiled_code, namespace),
|
||
{},
|
||
rule.timeout
|
||
)
|
||
|
||
# 获取入口函数
|
||
if rule.entry_function not in namespace:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=f"代码未定义指定的入口函数 '{rule.entry_function}'"
|
||
)
|
||
|
||
entry_func = namespace[rule.entry_function]
|
||
if not callable(entry_func):
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=f"'{rule.entry_function}' 不是一个可调用的函数"
|
||
)
|
||
|
||
# 调用入口函数
|
||
result = self._execute_with_timeout(
|
||
entry_func,
|
||
tuple(),
|
||
{},
|
||
rule.timeout
|
||
)
|
||
|
||
# 处理结果
|
||
if isinstance(result, dict):
|
||
# 如果函数返回一个字典，将它转换为ValidationResult
|
||
return ValidationResult(
|
||
is_valid=bool(result.get('is_valid', False)),
|
||
message=result.get('message'),
|
||
details=result.get('details', {})
|
||
)
|
||
else:
|
||
# 如果返回其他类型，将其解释为布尔值
|
||
return ValidationResult(
|
||
is_valid=bool(result),
|
||
message=str(result) if result is not True else "验证通过"
|
||
)
|
||
|
||
except TimeoutError as e:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=str(e),
|
||
exception=e
|
||
)
|
||
except Exception as e:
|
||
return ValidationResult(
|
||
is_valid=False,
|
||
message=f"执行代码时发生错误: {e}",
|
||
exception=e
|
||
) |