AIclinicalresearch/extraction_service/services/dc_executor.py

"""
DC工具C - Pandas代码执行服务

功能：
- AST静态代码检查（安全验证）
- Pandas代码执行（沙箱环境）
- 危险模块拦截
- 超时保护
"""

import ast
import sys
import pandas as pd
import numpy as np
import io
import traceback
from typing import Dict, Any, List, Tuple
from loguru import logger
import signal
import json


# ==================== 配置常量 ====================

# 危险模块黑名单（禁止导入）
DANGEROUS_MODULES = {
    'os', 'sys', 'subprocess', 'shutil', 'glob',
    'socket', 'urllib', 'requests', 'http',
    'pickle', 'shelve', 'dbm',
    'importlib', '__import__',
    'eval', 'exec', 'compile',
    'open', 'input', 'file',
}

# 危险内置函数黑名单
DANGEROUS_BUILTINS = {
    'eval', 'exec', 'compile', '__import__',
    'open', 'input', 'file',
    'getattr', 'setattr', 'delattr',
    'globals', 'locals', 'vars',
}

# 允许的安全模块
SAFE_MODULES = {
    'pandas', 'numpy', 'pd', 'np',
    'datetime', 'math', 'statistics',
    'json', 're', 'collections',
}

# 代码执行超时时间（秒）
EXECUTION_TIMEOUT = 30


# ==================== AST安全检查 ====================

class SecurityVisitor(ast.NodeVisitor):
    """
    AST访问器 - 检查代码中的危险操作
    """
    
    def __init__(self):
        self.errors: List[str] = []
        self.warnings: List[str] = []
    
    def visit_Import(self, node: ast.Import):
        """检查import语句"""
        for alias in node.names:
            module = alias.name.split('.')[0]
            if module in DANGEROUS_MODULES:
                self.errors.append(
                    f"🚫 禁止导入危险模块: {module} (行 {node.lineno})"
                )
            elif module not in SAFE_MODULES:
                self.warnings.append(
                    f"⚠️  导入了未知模块: {module} (行 {node.lineno})"
                )
        self.generic_visit(node)
    
    def visit_ImportFrom(self, node: ast.ImportFrom):
        """检查from...import语句"""
        if node.module:
            module = node.module.split('.')[0]
            if module in DANGEROUS_MODULES:
                self.errors.append(
                    f"🚫 禁止从危险模块导入: {module} (行 {node.lineno})"
                )
            elif module not in SAFE_MODULES:
                self.warnings.append(
                    f"⚠️  从未知模块导入: {module} (行 {node.lineno})"
                )
        self.generic_visit(node)
    
    def visit_Call(self, node: ast.Call):
        """检查函数调用"""
        # 检查是否调用了危险内置函数
        if isinstance(node.func, ast.Name):
            func_name = node.func.id
            if func_name in DANGEROUS_BUILTINS:
                self.errors.append(
                    f"🚫 禁止调用危险函数: {func_name}() (行 {node.lineno})"
                )
        
        # 检查是否调用了open文件操作
        if isinstance(node.func, ast.Attribute):
            if node.func.attr == 'open':
                self.errors.append(
                    f"🚫 禁止文件操作: open() (行 {node.lineno})"
                )
        
        self.generic_visit(node)
    
    def visit_Delete(self, node: ast.Delete):
        """检查删除操作"""
        # 不允许删除df变量本身
        for target in node.targets:
            if isinstance(target, ast.Name) and target.id == 'df':
                self.errors.append(
                    f"🚫 禁止删除DataFrame变量: del df (行 {node.lineno})"
                )
        self.generic_visit(node)


def validate_code(code: str) -> Dict[str, Any]:
    """
    AST静态代码检查
    
    Args:
        code: 待检查的Python代码
    
    Returns:
        {
            "valid": bool,
            "errors": List[str],
            "warnings": List[str]
        }
    """
    try:
        logger.info(f"开始AST代码检查，代码长度: {len(code)} 字符")
        
        # 解析AST
        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return {
                "valid": False,
                "errors": [f"❌ 语法错误 (行 {e.lineno}): {e.msg}"],
                "warnings": []
            }
        
        # 安全检查
        visitor = SecurityVisitor()
        visitor.visit(tree)
        
        # 额外检查：代码中是否包含df变量
        has_df = any(
            isinstance(node, ast.Name) and node.id == 'df'
            for node in ast.walk(tree)
        )
        
        if not has_df:
            visitor.warnings.append(
                "⚠️  代码中未使用 df 变量，可能无法操作数据"
            )
        
        # 返回结果
        is_valid = len(visitor.errors) == 0
        
        logger.info(
            f"AST检查完成: valid={is_valid}, "
            f"errors={len(visitor.errors)}, warnings={len(visitor.warnings)}"
        )
        
        return {
            "valid": is_valid,
            "errors": visitor.errors,
            "warnings": visitor.warnings
        }
    
    except Exception as e:
        logger.error(f"AST检查失败: {str(e)}")
        return {
            "valid": False,
            "errors": [f"❌ 检查失败: {str(e)}"],
            "warnings": []
        }


# ==================== 超时处理 ====================

class TimeoutException(Exception):
    """代码执行超时异常"""
    pass


def timeout_handler(signum, frame):
    """超时信号处理器"""
    raise TimeoutException("代码执行超时（>30秒）")


# ==================== Pandas代码执行 ====================

def execute_pandas_code(data: List[Dict[str, Any]], code: str) -> Dict[str, Any]:
    """
    在沙箱环境中执行Pandas代码
    
    Args:
        data: JSON格式的数据（数组对象）
        code: Pandas代码（必须操作df变量）
    
    Returns:
        {
            "success": bool,
            "result_data": List[Dict],  # 执行后的数据
            "output": str,              # 打印输出
            "error": str,               # 错误信息（如果失败）
            "execution_time": float     # 执行时间（秒）
        }
    """
    import time
    start_time = time.time()
    
    try:
        logger.info(f"开始执行Pandas代码，数据行数: {len(data)}")
        
        # 1. 先进行AST检查
        validation = validate_code(code)
        if not validation["valid"]:
            return {
                "success": False,
                "result_data": None,
                "output": "",
                "error": f"代码未通过安全检查:\n" + "\n".join(validation["errors"]),
                "execution_time": time.time() - start_time
            }
        
        # 2. 创建DataFrame
        try:
            df = pd.DataFrame(data)
            logger.info(f"DataFrame创建成功: shape={df.shape}")
        except Exception as e:
            return {
                "success": False,
                "result_data": None,
                "output": "",
                "error": f"数据转换失败: {str(e)}",
                "execution_time": time.time() - start_time
            }
        
        # 3. 准备安全的执行环境
        safe_globals = {
            'pd': pd,
            'np': np,
            'df': df,
            '__builtins__': {
                # 只允许安全的内置函数
                'len': len,
                'range': range,
                'enumerate': enumerate,
                'zip': zip,
                'map': map,
                'filter': filter,
                'list': list,
                'dict': dict,
                'set': set,
                'tuple': tuple,
                'str': str,
                'int': int,
                'float': float,
                'bool': bool,
                'print': print,
                'sum': sum,
                'min': min,
                'max': max,
                'abs': abs,
                'round': round,
                'sorted': sorted,
                'reversed': reversed,
                'any': any,
                'all': all,
            }
        }
        
        # 4. 捕获print输出
        old_stdout = sys.stdout
        sys.stdout = captured_output = io.StringIO()
        
        try:
            # 5. 设置超时保护（仅在Unix系统上）
            if hasattr(signal, 'SIGALRM'):
                signal.signal(signal.SIGALRM, timeout_handler)
                signal.alarm(EXECUTION_TIMEOUT)
            
            # 6. 执行代码
            logger.info("执行用户代码...")
            exec(code, safe_globals)
            
            # 7. 取消超时
            if hasattr(signal, 'SIGALRM'):
                signal.alarm(0)
            
            # 8. 获取执行后的DataFrame
            df_result = safe_globals['df']
            
            # 9. 验证结果
            if not isinstance(df_result, pd.DataFrame):
                raise ValueError(
                    f"执行后df不是DataFrame类型，而是 {type(df_result)}"
                )
            
            logger.info(f"代码执行成功: shape={df_result.shape}")
            
            # 10. 转换回JSON格式（NaN转为None避免JSON序列化错误）
            result_data = df_result.replace({np.nan: None}).to_dict('records')
            
            # 11. 获取print输出
            output = captured_output.getvalue()
            
            execution_time = time.time() - start_time
            logger.info(f"执行完成，耗时: {execution_time:.3f}秒")
            
            return {
                "success": True,
                "result_data": result_data,
                "output": output.strip() if output else "",
                "error": None,
                "execution_time": execution_time,
                "result_shape": df_result.shape
            }
        
        except TimeoutException as e:
            logger.error(f"代码执行超时: {str(e)}")
            return {
                "success": False,
                "result_data": None,
                "output": captured_output.getvalue(),
                "error": f"⏱️  执行超时: 代码运行超过 {EXECUTION_TIMEOUT} 秒",
                "execution_time": time.time() - start_time
            }
        
        except Exception as e:
            logger.error(f"代码执行失败: {str(e)}")
            error_traceback = traceback.format_exc()
            return {
                "success": False,
                "result_data": None,
                "output": captured_output.getvalue(),
                "error": f"❌ 执行错误:\n{str(e)}\n\n{error_traceback}",
                "execution_time": time.time() - start_time
            }
        
        finally:
            # 恢复stdout
            sys.stdout = old_stdout
            
            # 取消超时
            if hasattr(signal, 'SIGALRM'):
                signal.alarm(0)
    
    except Exception as e:
        logger.error(f"代码执行服务失败: {str(e)}")
        return {
            "success": False,
            "result_data": None,
            "output": "",
            "error": f"服务错误: {str(e)}",
            "execution_time": time.time() - start_time
        }


# ==================== 测试函数 ====================

def test_dc_executor():
    """测试DC执行器"""
    
    print("=" * 60)
    print("测试1: AST安全检查 - 正常代码")
    print("=" * 60)
    
    safe_code = """
import pandas as pd
df['age_group'] = df['age'].apply(lambda x: '老年' if x > 60 else '非老年')
print(df['age_group'].value_counts())
"""
    
    result = validate_code(safe_code)
    print(json.dumps(result, indent=2, ensure_ascii=False))
    
    print("\n" + "=" * 60)
    print("测试2: AST安全检查 - 危险代码")
    print("=" * 60)
    
    dangerous_code = """
import os
import sys
os.system('rm -rf /')
"""
    
    result = validate_code(dangerous_code)
    print(json.dumps(result, indent=2, ensure_ascii=False))
    
    print("\n" + "=" * 60)
    print("测试3: 代码执行 - 简单操作")
    print("=" * 60)
    
    test_data = [
        {"patient_id": "P001", "age": 25, "gender": "男"},
        {"patient_id": "P002", "age": 65, "gender": "女"},
        {"patient_id": "P003", "age": 45, "gender": "男"},
    ]
    
    simple_code = """
df['age_group'] = df['age'].apply(lambda x: '老年' if x > 60 else '非老年')
print(f"处理完成，共 {len(df)} 行")
"""
    
    result = execute_pandas_code(test_data, simple_code)
    print(json.dumps(result, indent=2, ensure_ascii=False))
    
    print("\n" + "=" * 60)
    print("测试完成！")
    print("=" * 60)


if __name__ == "__main__":
    test_dc_executor()