DinQuant/backend_api_python/app/utils/safe_exec.py

"""
Secure code execution tools
Provides timeouts, resource limits and a sandbox environment
"""
import signal
import sys
import os
import threading
import traceback
from typing import Dict, Any, Optional, Tuple
from contextlib import contextmanager

from app.utils.logger import get_logger

logger = get_logger(__name__)


class TimeoutError(Exception):
    """Code execution timeout exception"""
    pass


@contextmanager
def timeout_context(seconds: int):
    """
    Code execution timeout context manager
    
    Notice:
    - Only works on Unix/Linux systems
    - Only valid in the main thread, non-main threads will be downgraded to unlimited timeout
    - Will downgrade to unlimited timeout on Windows
    
    Args:
        seconds: timeout (seconds)
    """
    # Check if in main thread
    is_main_thread = threading.current_thread() is threading.main_thread()
    
    if sys.platform == 'win32':
        # signal.alarm is not supported on Windows, only warnings can be logged
        logger.warning("Windows does not support signal-based timeouts; execution time limits may not work")
        yield
        return
    
    if not is_main_thread:
        # Non-main threads cannot use signal. Warnings are logged but timeouts are not limited.
        # logger.warning(f"Currently running in a non-main thread (thread: {threading.current_thread().name}),"
        #               f"signal timeout is not available, code execution may not limit the time")
        yield
        return
    
    def timeout_handler(signum, frame):
        raise TimeoutError(f"Code execution timed out after {seconds} seconds")
    
    try:
        # Set up signal handler
        old_handler = signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(seconds)
        
        try:
            yield
        finally:
            # Restore original signal handler
            signal.alarm(0)
            signal.signal(signal.SIGALRM, old_handler)
    except ValueError as e:
        # If signal setup fails (such as in some environments), log a warning but do not interrupt execution
        logger.warning(f"Failed to set signal timeout: {str(e)}; execution will continue without timeout enforcement")
        yield


def safe_exec_code(
    code: str,
    exec_globals: Dict[str, Any],
    exec_locals: Optional[Dict[str, Any]] = None,
    timeout: int = 30,
    max_memory_mb: Optional[int] = None
) -> Dict[str, Any]:
    """
    Safe execution of Python code
    
    Args:
        code: Python code to execute
        exec_globals: dictionary of global variables
        exec_locals: dictionary of local variables (if None, exec_globals is used)
        timeout: timeout (seconds), default 30 seconds
        max_memory_mb: Maximum memory limit (MB), default 500MB
    
    Returns:
        Execution result dictionary, including:
        - success: bool, whether the execution was successful
        - error: str, error message (if failure)
        - result: Any, execution result (if any)
    
    Raises:
        TimeoutError: If the code execution times out
    """
    if exec_locals is None:
        exec_locals = exec_globals
    
    # Set memory limit (if supported)
    if max_memory_mb is None:
        max_memory_mb = 500  # Default 500MB
    
    try:
        # Note: resource.setrlimit is process level and affects the entire API process.
        # The previous global limit of 500MB could prevent parallel strategies/threads from allocating memory.
        # Only set if SAFE_EXEC_ENABLE_RLIMIT is explicitly turned on.
        if sys.platform != 'win32' and os.getenv('SAFE_EXEC_ENABLE_RLIMIT', 'false').lower() == 'true':
            try:
                import resource
                max_memory_bytes = max_memory_mb * 1024 * 1024
                resource.setrlimit(resource.RLIMIT_AS, (max_memory_bytes, max_memory_bytes))
                logger.debug(f"Memory limit set: {max_memory_mb}MB (SAFE_EXEC_ENABLE_RLIMIT enabled)")
            except (ImportError, ValueError, OSError) as e:
                logger.warning(f"Failed to set memory limit: {str(e)}")
        else:
            logger.debug("No resource memory limit (SAFE_EXEC_ENABLE_RLIMIT disabled or unsupported platform)")
        
        # On Windows, timeout_context doesn't really limit the time
        # But a warning will be logged
        with timeout_context(timeout):
            exec(code, exec_globals, exec_locals)
        
        return {
            'success': True,
            'error': None,
            'result': None
        }
        
    except MemoryError as e:
        error_msg = f"Code execution exceeded the memory limit of {max_memory_mb}MB"
        logger.error(f"Code execution out of memory (limit={max_memory_mb}MB)")
        return {
            'success': False,
            'error': error_msg,
            'result': None
        }
    except TimeoutError as e:
        error_msg = str(e)
        logger.error(f"Code execution timed out (timeout={timeout}s)")
        return {
            'success': False,
            'error': error_msg,
            'result': None
        }
    except Exception as e:
        error_msg = f"Code execution error: {str(e)}\n{traceback.format_exc()}"
        logger.error(f"Code execution error: {str(e)}")
        logger.error(traceback.format_exc())
        return {
            'success': False,
            'error': error_msg,
            'result': None
        }


def validate_code_safety(code: str) -> Tuple[bool, Optional[str]]:
    """
    Verify code security (basic check)
    
    Check your code for dangerous function calls or imports
    
    Args:
        code: Python code to check
    
    Returns:
        (is_safe: bool, error_message: Optional[str])
    """
    import ast
    import re
    
    # Dangerous keywords and function names
    dangerous_patterns = [
        # System command execution
        r'\bos\.system\b',
        r'\bos\.popen\b',
        r'\bos\.spawn\b',
        r'\bos\.exec\b',
        r'\bos\.fork\b',
        r'\bsubprocess\b',
        r'\bcommands\b',
        # code execution
        r'\b__import__\s*\(',
        r'\beval\s*\(',
        r'\bexec\s*\(',
        r'\bcompile\s*\(',
        # File operations
        r'\bopen\s*\(',
        r'\bfile\s*\(',
        r'\b__builtins__\b',
        # module import
        r'\bimport\s+os\b',
        r'\bimport\s+sys\b',
        r'\bimport\s+subprocess\b',
        r'\bimport\s+pymysql\b',
        r'\bimport\s+sqlite3\b',
        r'\bimport\s+requests\b',
        r'\bimport\s+urllib\b',
        r'\bimport\s+http\b',
        r'\bimport\s+socket\b',
        r'\bimport\s+ftplib\b',
        r'\bimport\s+telnetlib\b',
        r'\bimport\s+pickle\b',
        r'\bimport\s+cpickle\b',
        r'\bimport\s+marshal\b',
        r'\bimport\s+ctypes\b',
        r'\bimport\s+multiprocessing\b',
        r'\bimport\s+threading\b',
        r'\bimport\s+concurrent\b',
        # Reflection and metaprogramming (possibly used to bypass restrictions)
        r'\bgetattr\s*\(.*__import__',
        r'\bgetattr\s*\(.*eval',
        r'\bgetattr\s*\(.*exec',
        r'\bsetattr\s*\(',
        r'\b__getattribute__\b',
        r'\b__setattr__\b',
        r'\b__dict__\b',
        r'\bglobals\s*\(',
        r'\blocals\s*\(',
        r'\bdir\s*\(',
        r'\btype\s*\(.*\)\s*\(',  # type() may be used to create new types
        r'\b__class__\b',
        r'\b__bases__\b',
        r'\b__subclasses__\b',
        r'\b__mro__\b',
        r'\b__init__\b.*__import__',
        r'\b__new__\b.*__import__',
        # Other dangerous operations
        r'\b__builtins__\s*\[',
        r'\b__builtins__\s*\.',
        r'\b__import__\s*\(',
        r'\bimportlib\b',
        r'\bimp\b',
    ]
    
    # Check your code for dangerous patterns
    for pattern in dangerous_patterns:
        if re.search(pattern, code):
            return False, f"Dangerous code pattern detected: {pattern}"
    
    # Try parsing the AST, checking for dangerous nodes
    try:
        tree = ast.parse(code)
        
        # List of dangerous modules (extended)
        dangerous_modules = [
            'os', 'sys', 'subprocess', 'pymysql', 'sqlite3',
            'requests', 'urllib', 'http', 'socket', 'ftplib', 'telnetlib',
            'pickle', 'cpickle', 'marshal', 'ctypes',
            'multiprocessing', 'threading', 'concurrent',
            'importlib', 'imp', 'builtins'
        ]
        
        # List of dangerous functions (extended)
        # Note: hasattr is safe and is only used to check attributes, not to access
        dangerous_functions = [
            'eval', 'exec', 'compile', '__import__',
            'getattr', 'setattr', 'delattr',  # hasattr has been removed, it is safe
            'globals', 'locals', 'vars', 'dir', 'type'
        ]
        
        # Check for dangerous function calls
        for node in ast.walk(tree):
            # Check for calls to dangerous functions
            if isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name):
                    func_name = node.func.id
                    if func_name in dangerous_functions:
                        return False, f"Dangerous function call detected: {func_name}()"
                
                # Check if there are calls to os.system etc.
                if isinstance(node.func, ast.Attribute):
                    if isinstance(node.func.value, ast.Name):
                        if node.func.value.id in dangerous_modules:
                            return False, f"Dangerous module call detected: {node.func.value.id}.{node.func.attr}"
                    
                    # Check if there are bypass methods such as getattr(builtins, '__import__')
                    if isinstance(node.func, ast.Name) and node.func.id == 'getattr':
                        # Check the parameters of getattr
                        if len(node.args) >= 2:
                            if isinstance(node.args[0], ast.Name) and node.args[0].id in ['builtins', '__builtins__']:
                                if isinstance(node.args[1], ast.Constant) and node.args[1].value in dangerous_functions:
                                    return False, f"Detected an attempt to bypass restrictions via getattr: getattr({node.args[0].id}, '{node.args[1].value}')"
        
        # Check the import statement: the use of import is prohibited in user scripts (security dependencies are uniformly injected by the platform)
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                return False, "Import statements are not allowed in scripts. Use the platform-provided objects such as pd and np directly."
            
            if isinstance(node, ast.ImportFrom):
                return False, "Import statements are not allowed in scripts. Use the platform-provided objects such as pd and np directly."
        
        # Check if there is an attempt to access __builtins__
        for node in ast.walk(tree):
            if isinstance(node, ast.Attribute):
                if isinstance(node.attr, str) and node.attr.startswith('__') and node.attr.endswith('__'):
                    if node.attr in ['__builtins__', '__import__', '__class__', '__bases__', '__subclasses__', '__mro__']:
                        # Check if used in dangerous context
                        if isinstance(node.value, ast.Name) and node.value.id in ['builtins', '__builtins__']:
                            return False, f"Dangerous attribute access detected: {node.value.id}.{node.attr}"
        
    except SyntaxError as e:
        return False, f"Code syntax error: {str(e)}"
    except Exception as e:
        # If AST parsing fails, log a warning but allow to continue (possibly incomplete code)
        logger.warning(f"AST parse failed; skipping safety checks: {str(e)}")
    
    return True, None