Refactor code for improved readability and consistency
- Cleaned up whitespace and formatting in various files including http.py, language.py, logger.py, safe_exec.py, and SQL migration scripts. - Consolidated import statements and removed unnecessary blank lines. - Updated logging configuration for better clarity. - Enhanced the safe execution code with improved error handling and logging. - Removed commented-out code and unnecessary variables in backfill_zero_trades.py and other scripts. - Added a pyproject.toml for Ruff and Vulture configuration. - Introduced requirements-dev.txt for development dependencies. - Removed commented-out stock entries in init.sql for cleaner migration scripts.
This commit is contained in:
dienakdz
@@ -2,13 +2,14 @@
|
||||
Secure code execution tools
|
||||
Provides timeouts, resource limits and a sandbox environment
|
||||
"""
|
||||
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import os
|
||||
import threading
|
||||
import traceback
|
||||
from typing import Dict, Any, Optional, Tuple
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
@@ -17,6 +18,7 @@ logger = get_logger(__name__)
|
||||
|
||||
class TimeoutError(Exception):
|
||||
"""Code execution timeout exception"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@@ -24,39 +26,39 @@ class TimeoutError(Exception):
|
||||
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':
|
||||
|
||||
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:
|
||||
@@ -74,41 +76,42 @@ def safe_exec_code(
|
||||
exec_globals: Dict[str, Any],
|
||||
exec_locals: Optional[Dict[str, Any]] = None,
|
||||
timeout: int = 30,
|
||||
max_memory_mb: Optional[int] = None
|
||||
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':
|
||||
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)")
|
||||
@@ -116,150 +119,159 @@ def safe_exec_code(
|
||||
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:
|
||||
|
||||
return {"success": True, "error": None, "result": None}
|
||||
|
||||
except MemoryError:
|
||||
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
|
||||
}
|
||||
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
|
||||
}
|
||||
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
|
||||
}
|
||||
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',
|
||||
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*\(',
|
||||
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',
|
||||
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',
|
||||
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__',
|
||||
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',
|
||||
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'
|
||||
"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'
|
||||
"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
|
||||
@@ -268,42 +280,58 @@ def validate_code_safety(code: str) -> Tuple[bool, Optional[str]]:
|
||||
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':
|
||||
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[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}')"
|
||||
|
||||
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."
|
||||
|
||||
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."
|
||||
|
||||
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__']:
|
||||
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__']:
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user