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
@@ -1,9 +1,9 @@
|
||||
"""
|
||||
tool module
|
||||
"""
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
from app.utils.cache import CacheManager
|
||||
from app.utils.http import get_retry_session
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
__all__ = ['get_logger', 'CacheManager', 'get_retry_session']
|
||||
|
||||
__all__ = ["get_logger", "CacheManager", "get_retry_session"]
|
||||
|
||||
@@ -4,44 +4,43 @@ Authentication Utilities
|
||||
JWT token generation, verification, and middleware decorators.
|
||||
Supports multi-user authentication with role-based access control.
|
||||
"""
|
||||
import jwt
|
||||
|
||||
import datetime
|
||||
import os
|
||||
from functools import wraps
|
||||
from flask import request, jsonify, g
|
||||
|
||||
import jwt
|
||||
from flask import g, jsonify, request
|
||||
|
||||
from app.config.settings import Config
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def generate_token(user_id: int, username: str, role: str = 'user', token_version: int = 1) -> str:
|
||||
def generate_token(user_id: int, username: str, role: str = "user", token_version: int = 1) -> str:
|
||||
"""
|
||||
Generate JWT token with user information.
|
||||
|
||||
|
||||
Args:
|
||||
user_id: User ID
|
||||
username: Username
|
||||
role: User role (admin/manager/user/viewer)
|
||||
token_version: Token version for single-client enforcement
|
||||
|
||||
|
||||
Returns:
|
||||
JWT token string
|
||||
"""
|
||||
try:
|
||||
payload = {
|
||||
'exp': datetime.datetime.utcnow() + datetime.timedelta(days=7),
|
||||
'iat': datetime.datetime.utcnow(),
|
||||
'sub': username,
|
||||
'user_id': user_id,
|
||||
'role': role,
|
||||
'token_version': token_version, # For single client login control
|
||||
"exp": datetime.datetime.utcnow() + datetime.timedelta(days=7),
|
||||
"iat": datetime.datetime.utcnow(),
|
||||
"sub": username,
|
||||
"user_id": user_id,
|
||||
"role": role,
|
||||
"token_version": token_version, # For single client login control
|
||||
}
|
||||
return jwt.encode(
|
||||
payload,
|
||||
Config.SECRET_KEY,
|
||||
algorithm='HS256'
|
||||
)
|
||||
return jwt.encode(payload, Config.SECRET_KEY, algorithm="HS256")
|
||||
except Exception as e:
|
||||
logger.error(f"Token generation failed: {e}")
|
||||
return None
|
||||
@@ -50,26 +49,26 @@ def generate_token(user_id: int, username: str, role: str = 'user', token_versio
|
||||
def verify_token(token: str) -> dict:
|
||||
"""
|
||||
Verify JWT token and return payload.
|
||||
|
||||
|
||||
Args:
|
||||
token: JWT token string
|
||||
|
||||
|
||||
Returns:
|
||||
Token payload dict or None if invalid
|
||||
"""
|
||||
try:
|
||||
payload = jwt.decode(token, Config.SECRET_KEY, algorithms=['HS256'])
|
||||
|
||||
payload = jwt.decode(token, Config.SECRET_KEY, algorithms=["HS256"])
|
||||
|
||||
# Verify token_version (single client login control)
|
||||
user_id = payload.get('user_id')
|
||||
token_version = payload.get('token_version')
|
||||
|
||||
user_id = payload.get("user_id")
|
||||
token_version = payload.get("token_version")
|
||||
|
||||
if user_id and token_version is not None:
|
||||
# Check if token_version in database matches
|
||||
if not _verify_token_version(user_id, token_version):
|
||||
logger.debug(f"Token version mismatch for user {user_id}: expected current, got {token_version}")
|
||||
return None
|
||||
|
||||
|
||||
return payload
|
||||
except jwt.ExpiredSignatureError:
|
||||
logger.debug("Token expired")
|
||||
@@ -83,29 +82,27 @@ def _verify_token_version(user_id: int, token_version: int) -> bool:
|
||||
"""
|
||||
Verify that the token version matches the version stored in the database.
|
||||
Used to implement single client login (kick out duplicate logins).
|
||||
|
||||
|
||||
Args:
|
||||
user_id: user ID
|
||||
token_version: version number in Token
|
||||
|
||||
|
||||
Returns:
|
||||
True if version matches, False otherwise
|
||||
"""
|
||||
try:
|
||||
from app.utils.db import get_db_connection
|
||||
|
||||
with get_db_connection() as db:
|
||||
cur = db.cursor()
|
||||
cur.execute(
|
||||
"SELECT token_version FROM qd_users WHERE id = ?",
|
||||
(user_id,)
|
||||
)
|
||||
cur.execute("SELECT token_version FROM qd_users WHERE id = ?", (user_id,))
|
||||
row = cur.fetchone()
|
||||
cur.close()
|
||||
|
||||
|
||||
if not row:
|
||||
return False
|
||||
|
||||
db_token_version = row.get('token_version') or 1
|
||||
|
||||
db_token_version = row.get("token_version") or 1
|
||||
return int(token_version) == int(db_token_version)
|
||||
except Exception as e:
|
||||
logger.error(f"_verify_token_version failed: {e}")
|
||||
@@ -115,45 +112,46 @@ def _verify_token_version(user_id: int, token_version: int) -> bool:
|
||||
|
||||
def get_current_user_id() -> int:
|
||||
"""Get current user ID from flask.g context"""
|
||||
return getattr(g, 'user_id', None)
|
||||
return getattr(g, "user_id", None)
|
||||
|
||||
|
||||
def get_current_user_role() -> str:
|
||||
"""Get current user role from flask.g context"""
|
||||
return getattr(g, 'user_role', 'user')
|
||||
return getattr(g, "user_role", "user")
|
||||
|
||||
|
||||
def login_required(f):
|
||||
"""
|
||||
Decorator that enforces Bearer token auth.
|
||||
|
||||
|
||||
Sets g.user, g.user_id, g.user_role on successful auth.
|
||||
"""
|
||||
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
token = None
|
||||
|
||||
|
||||
# Read token from Authorization: Bearer <token>
|
||||
auth_header = request.headers.get('Authorization')
|
||||
auth_header = request.headers.get("Authorization")
|
||||
if auth_header:
|
||||
parts = auth_header.split()
|
||||
if len(parts) == 2 and parts[0].lower() == 'bearer':
|
||||
if len(parts) == 2 and parts[0].lower() == "bearer":
|
||||
token = parts[1]
|
||||
|
||||
|
||||
if not token:
|
||||
return jsonify({'code': 401, 'msg': 'Token missing', 'data': None}), 401
|
||||
|
||||
return jsonify({"code": 401, "msg": "Token missing", "data": None}), 401
|
||||
|
||||
payload = verify_token(token)
|
||||
if not payload:
|
||||
return jsonify({'code': 401, 'msg': 'Token invalid or expired', 'data': None}), 401
|
||||
|
||||
return jsonify({"code": 401, "msg": "Token invalid or expired", "data": None}), 401
|
||||
|
||||
# Store user info in flask.g
|
||||
g.user = payload.get('sub')
|
||||
g.user_id = payload.get('user_id')
|
||||
g.user_role = payload.get('role', 'user')
|
||||
|
||||
g.user = payload.get("sub")
|
||||
g.user_id = payload.get("user_id")
|
||||
g.user_role = payload.get("role", "user")
|
||||
|
||||
return f(*args, **kwargs)
|
||||
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
@@ -162,12 +160,14 @@ def admin_required(f):
|
||||
Decorator that requires admin role.
|
||||
Must be used after @login_required.
|
||||
"""
|
||||
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
role = getattr(g, 'user_role', None)
|
||||
if role != 'admin':
|
||||
return jsonify({'code': 403, 'msg': 'Admin access required', 'data': None}), 403
|
||||
role = getattr(g, "user_role", None)
|
||||
if role != "admin":
|
||||
return jsonify({"code": 403, "msg": "Admin access required", "data": None}), 403
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
@@ -176,12 +176,14 @@ def manager_required(f):
|
||||
Decorator that requires manager or admin role.
|
||||
Must be used after @login_required.
|
||||
"""
|
||||
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
role = getattr(g, 'user_role', None)
|
||||
if role not in ('admin', 'manager'):
|
||||
return jsonify({'code': 403, 'msg': 'Manager access required', 'data': None}), 403
|
||||
role = getattr(g, "user_role", None)
|
||||
if role not in ("admin", "manager"):
|
||||
return jsonify({"code": 403, "msg": "Manager access required", "data": None}), 403
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
@@ -189,38 +191,38 @@ def permission_required(permission: str):
|
||||
"""
|
||||
Decorator factory that checks for a specific permission.
|
||||
Must be used after @login_required.
|
||||
|
||||
|
||||
Usage:
|
||||
@login_required
|
||||
@permission_required('strategy')
|
||||
def my_endpoint():
|
||||
...
|
||||
"""
|
||||
|
||||
def decorator(f):
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
role = getattr(g, 'user_role', 'user')
|
||||
|
||||
role = getattr(g, "user_role", "user")
|
||||
|
||||
# Import here to avoid circular import
|
||||
from app.services.user_service import get_user_service
|
||||
|
||||
permissions = get_user_service().get_user_permissions(role)
|
||||
|
||||
|
||||
if permission not in permissions:
|
||||
return jsonify({
|
||||
'code': 403,
|
||||
'msg': f'Permission denied: {permission}',
|
||||
'data': None
|
||||
}), 403
|
||||
|
||||
return jsonify({"code": 403, "msg": f"Permission denied: {permission}", "data": None}), 403
|
||||
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
# Legacy compatibility: single-user mode fallback
|
||||
def _is_single_user_mode() -> bool:
|
||||
"""Check if system is in single-user (legacy) mode"""
|
||||
return os.getenv('SINGLE_USER_MODE', 'false').lower() == 'true'
|
||||
return os.getenv("SINGLE_USER_MODE", "false").lower() == "true"
|
||||
|
||||
|
||||
def authenticate_legacy(username: str, password: str) -> dict:
|
||||
@@ -230,9 +232,9 @@ def authenticate_legacy(username: str, password: str) -> dict:
|
||||
"""
|
||||
if username == Config.ADMIN_USER and password == Config.ADMIN_PASSWORD:
|
||||
return {
|
||||
'user_id': 1,
|
||||
'username': username,
|
||||
'role': 'admin',
|
||||
'nickname': 'Admin',
|
||||
"user_id": 1,
|
||||
"username": username,
|
||||
"role": "admin",
|
||||
"nickname": "Admin",
|
||||
}
|
||||
return None
|
||||
|
||||
@@ -3,24 +3,25 @@ Cache utilities.
|
||||
Local-first behavior: use in-memory cache by default.
|
||||
Redis is only used when explicitly enabled via environment variables.
|
||||
"""
|
||||
import time
|
||||
import threading
|
||||
from typing import Optional, Any
|
||||
import json
|
||||
|
||||
from app.utils.logger import get_logger
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Optional
|
||||
|
||||
from app.config import CacheConfig
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class MemoryCache:
|
||||
"""In-memory cache (an alternative if Redis is unavailable)"""
|
||||
|
||||
|
||||
def __init__(self):
|
||||
self._cache = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
|
||||
def get(self, key: str) -> Optional[str]:
|
||||
with self._lock:
|
||||
if key in self._cache:
|
||||
@@ -30,17 +31,17 @@ class MemoryCache:
|
||||
else:
|
||||
del self._cache[key]
|
||||
return None
|
||||
|
||||
|
||||
def setex(self, key: str, ttl: int, value: str):
|
||||
with self._lock:
|
||||
expiry = time.time() + ttl
|
||||
self._cache[key] = (value, expiry)
|
||||
|
||||
|
||||
def delete(self, key: str):
|
||||
with self._lock:
|
||||
if key in self._cache:
|
||||
del self._cache[key]
|
||||
|
||||
|
||||
def clear(self):
|
||||
with self._lock:
|
||||
self._cache.clear()
|
||||
@@ -48,10 +49,10 @@ class MemoryCache:
|
||||
|
||||
class CacheManager:
|
||||
"""cache manager"""
|
||||
|
||||
|
||||
_instance = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
@@ -59,11 +60,11 @@ class CacheManager:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._initialized = False
|
||||
return cls._instance
|
||||
|
||||
|
||||
def __init__(self):
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
|
||||
self._initialized = True
|
||||
self._client = None
|
||||
self._use_redis = False
|
||||
@@ -77,6 +78,7 @@ class CacheManager:
|
||||
# Try Redis only when enabled.
|
||||
try:
|
||||
import redis
|
||||
|
||||
from app.config import RedisConfig
|
||||
|
||||
self._client = redis.Redis(
|
||||
@@ -86,7 +88,7 @@ class CacheManager:
|
||||
password=RedisConfig.PASSWORD,
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=RedisConfig.CONNECT_TIMEOUT,
|
||||
socket_timeout=RedisConfig.SOCKET_TIMEOUT
|
||||
socket_timeout=RedisConfig.SOCKET_TIMEOUT,
|
||||
)
|
||||
self._client.ping()
|
||||
self._use_redis = True
|
||||
@@ -96,7 +98,7 @@ class CacheManager:
|
||||
logger.info(f"Redis is enabled but unavailable; using in-memory cache instead: {e}")
|
||||
self._client = MemoryCache()
|
||||
self._use_redis = False
|
||||
|
||||
|
||||
def get(self, key: str) -> Optional[Any]:
|
||||
"""Get cache"""
|
||||
try:
|
||||
@@ -107,22 +109,21 @@ class CacheManager:
|
||||
except Exception as e:
|
||||
logger.error(f"Cache read failed: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def set(self, key: str, value: Any, ttl: int = 300):
|
||||
"""Set up cache"""
|
||||
try:
|
||||
self._client.setex(key, ttl, json.dumps(value))
|
||||
except Exception as e:
|
||||
logger.error(f"Cache write failed: {e}")
|
||||
|
||||
|
||||
def delete(self, key: str):
|
||||
"""Delete cache"""
|
||||
try:
|
||||
self._client.delete(key)
|
||||
except Exception as e:
|
||||
logger.error(f"Cache delete failed: {e}")
|
||||
|
||||
|
||||
@property
|
||||
def is_redis(self) -> bool:
|
||||
return self._use_redis
|
||||
|
||||
|
||||
@@ -10,8 +10,9 @@ flat keys like `openrouter.api_key` become nested dicts like:
|
||||
"openrouter": {"api_key": "..."}
|
||||
}
|
||||
"""
|
||||
from typing import Dict, Any, Optional, List, Tuple
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
@@ -31,15 +32,15 @@ def load_addon_config() -> Dict[str, Any]:
|
||||
Nested config dict (PHP-compatible shape)
|
||||
"""
|
||||
global _config_cache
|
||||
|
||||
|
||||
# If the cache exists, return directly
|
||||
if _config_cache is not None:
|
||||
return _config_cache
|
||||
|
||||
|
||||
config: Dict[str, Any] = {}
|
||||
|
||||
def set_nested(cfg: Dict[str, Any], dotted_key: str, value: Any) -> None:
|
||||
keys = dotted_key.split('.')
|
||||
keys = dotted_key.split(".")
|
||||
ref = cfg
|
||||
for i, k in enumerate(keys):
|
||||
if i == len(keys) - 1:
|
||||
@@ -54,82 +55,66 @@ def load_addon_config() -> Dict[str, Any]:
|
||||
if val is None:
|
||||
return None
|
||||
val = str(val).strip()
|
||||
return val if val != '' else None
|
||||
return val if val != "" else None
|
||||
|
||||
# Map env vars to PHP-style dotted keys.
|
||||
mappings: List[Tuple[str, str, str]] = [
|
||||
# internal
|
||||
('INTERNAL_API_KEY', 'internal_api.key', 'string'),
|
||||
|
||||
("INTERNAL_API_KEY", "internal_api.key", "string"),
|
||||
# OpenRouter / LLM
|
||||
('OPENROUTER_API_KEY', 'openrouter.api_key', 'string'),
|
||||
('OPENROUTER_API_URL', 'openrouter.api_url', 'string'),
|
||||
('OPENROUTER_MODEL', 'openrouter.model', 'string'),
|
||||
('OPENROUTER_TEMPERATURE', 'openrouter.temperature', 'float'),
|
||||
('OPENROUTER_MAX_TOKENS', 'openrouter.max_tokens', 'int'),
|
||||
('OPENROUTER_TIMEOUT', 'openrouter.timeout', 'int'),
|
||||
('OPENROUTER_CONNECT_TIMEOUT', 'openrouter.connect_timeout', 'int'),
|
||||
|
||||
("OPENROUTER_API_KEY", "openrouter.api_key", "string"),
|
||||
("OPENROUTER_API_URL", "openrouter.api_url", "string"),
|
||||
("OPENROUTER_MODEL", "openrouter.model", "string"),
|
||||
("OPENROUTER_TEMPERATURE", "openrouter.temperature", "float"),
|
||||
("OPENROUTER_MAX_TOKENS", "openrouter.max_tokens", "int"),
|
||||
("OPENROUTER_TIMEOUT", "openrouter.timeout", "int"),
|
||||
("OPENROUTER_CONNECT_TIMEOUT", "openrouter.connect_timeout", "int"),
|
||||
# OpenAI Direct
|
||||
('OPENAI_API_KEY', 'openai.api_key', 'string'),
|
||||
('OPENAI_BASE_URL', 'openai.base_url', 'string'),
|
||||
('OPENAI_MODEL', 'openai.model', 'string'),
|
||||
|
||||
("OPENAI_API_KEY", "openai.api_key", "string"),
|
||||
("OPENAI_BASE_URL", "openai.base_url", "string"),
|
||||
("OPENAI_MODEL", "openai.model", "string"),
|
||||
# Google Gemini
|
||||
('GOOGLE_API_KEY', 'google.api_key', 'string'),
|
||||
('GOOGLE_MODEL', 'google.model', 'string'),
|
||||
|
||||
("GOOGLE_API_KEY", "google.api_key", "string"),
|
||||
("GOOGLE_MODEL", "google.model", "string"),
|
||||
# DeepSeek
|
||||
('DEEPSEEK_API_KEY', 'deepseek.api_key', 'string'),
|
||||
('DEEPSEEK_BASE_URL', 'deepseek.base_url', 'string'),
|
||||
('DEEPSEEK_MODEL', 'deepseek.model', 'string'),
|
||||
|
||||
("DEEPSEEK_API_KEY", "deepseek.api_key", "string"),
|
||||
("DEEPSEEK_BASE_URL", "deepseek.base_url", "string"),
|
||||
("DEEPSEEK_MODEL", "deepseek.model", "string"),
|
||||
# xAI Grok
|
||||
('GROK_API_KEY', 'grok.api_key', 'string'),
|
||||
('GROK_BASE_URL', 'grok.base_url', 'string'),
|
||||
('GROK_MODEL', 'grok.model', 'string'),
|
||||
|
||||
("GROK_API_KEY", "grok.api_key", "string"),
|
||||
("GROK_BASE_URL", "grok.base_url", "string"),
|
||||
("GROK_MODEL", "grok.model", "string"),
|
||||
# LLM Provider Selection
|
||||
('LLM_PROVIDER', 'llm.provider', 'string'),
|
||||
|
||||
("LLM_PROVIDER", "llm.provider", "string"),
|
||||
# App
|
||||
('RATE_LIMIT', 'app.rate_limit', 'int'),
|
||||
('ENABLE_CACHE', 'app.enable_cache', 'bool'),
|
||||
('ENABLE_REQUEST_LOG', 'app.enable_request_log', 'bool'),
|
||||
|
||||
("RATE_LIMIT", "app.rate_limit", "int"),
|
||||
("ENABLE_CACHE", "app.enable_cache", "bool"),
|
||||
("ENABLE_REQUEST_LOG", "app.enable_request_log", "bool"),
|
||||
# Data source common
|
||||
('DATA_SOURCE_TIMEOUT', 'data_source.timeout', 'int'),
|
||||
('DATA_SOURCE_RETRY', 'data_source.retry_count', 'int'),
|
||||
('DATA_SOURCE_RETRY_BACKOFF', 'data_source.retry_backoff', 'float'),
|
||||
|
||||
("DATA_SOURCE_TIMEOUT", "data_source.timeout", "int"),
|
||||
("DATA_SOURCE_RETRY", "data_source.retry_count", "int"),
|
||||
("DATA_SOURCE_RETRY_BACKOFF", "data_source.retry_backoff", "float"),
|
||||
# Finnhub
|
||||
('FINNHUB_API_KEY', 'finnhub.api_key', 'string'),
|
||||
('FINNHUB_TIMEOUT', 'finnhub.timeout', 'int'),
|
||||
('FINNHUB_RATE_LIMIT', 'finnhub.rate_limit', 'int'),
|
||||
|
||||
("FINNHUB_API_KEY", "finnhub.api_key", "string"),
|
||||
("FINNHUB_TIMEOUT", "finnhub.timeout", "int"),
|
||||
("FINNHUB_RATE_LIMIT", "finnhub.rate_limit", "int"),
|
||||
# CCXT
|
||||
('CCXT_DEFAULT_EXCHANGE', 'ccxt.default_exchange', 'string'),
|
||||
('CCXT_TIMEOUT', 'ccxt.timeout', 'int'),
|
||||
|
||||
("CCXT_DEFAULT_EXCHANGE", "ccxt.default_exchange", "string"),
|
||||
("CCXT_TIMEOUT", "ccxt.timeout", "int"),
|
||||
# Other sources
|
||||
('YFINANCE_TIMEOUT', 'yfinance.timeout', 'int'),
|
||||
('AKSHARE_TIMEOUT', 'akshare.timeout', 'int'),
|
||||
('TIINGO_API_KEY', 'tiingo.api_key', 'string'),
|
||||
('TIINGO_TIMEOUT', 'tiingo.timeout', 'int'),
|
||||
('TWELVE_DATA_API_KEY', 'twelve_data.api_key', 'string'),
|
||||
|
||||
("YFINANCE_TIMEOUT", "yfinance.timeout", "int"),
|
||||
("TIINGO_API_KEY", "tiingo.api_key", "string"),
|
||||
("TIINGO_TIMEOUT", "tiingo.timeout", "int"),
|
||||
# Search (Google CSE / Bing)
|
||||
('SEARCH_PROVIDER', 'search.provider', 'string'),
|
||||
('SEARCH_MAX_RESULTS', 'search.max_results', 'int'),
|
||||
('SEARCH_GOOGLE_API_KEY', 'search.google.api_key', 'string'),
|
||||
('SEARCH_GOOGLE_CX', 'search.google.cx', 'string'),
|
||||
('SEARCH_BING_API_KEY', 'search.bing.api_key', 'string'),
|
||||
|
||||
("SEARCH_PROVIDER", "search.provider", "string"),
|
||||
("SEARCH_MAX_RESULTS", "search.max_results", "int"),
|
||||
("SEARCH_GOOGLE_API_KEY", "search.google.api_key", "string"),
|
||||
("SEARCH_GOOGLE_CX", "search.google.cx", "string"),
|
||||
("SEARCH_BING_API_KEY", "search.bing.api_key", "string"),
|
||||
# Tavily (AI-optimized search)
|
||||
('TAVILY_API_KEYS', 'tavily.api_keys', 'string'),
|
||||
|
||||
("TAVILY_API_KEYS", "tavily.api_keys", "string"),
|
||||
# SerpAPI (Google/Bing scraper)
|
||||
('SERPAPI_KEYS', 'serpapi.api_keys', 'string'),
|
||||
("SERPAPI_KEYS", "serpapi.api_keys", "string"),
|
||||
]
|
||||
|
||||
for env_name, dotted_key, value_type in mappings:
|
||||
@@ -149,77 +134,78 @@ def load_addon_config() -> Dict[str, Any]:
|
||||
def _convert_config_value(value: str, value_type: str) -> Any:
|
||||
"""
|
||||
Convert configuration values according to type (consistent with the convertConfigValue method on the PHP side)
|
||||
|
||||
|
||||
Args:
|
||||
value: configuration value string (may be None)
|
||||
value_type: configuration type
|
||||
|
||||
|
||||
Returns:
|
||||
Converted configuration value
|
||||
"""
|
||||
# Handling None or null values
|
||||
if value is None or value == '':
|
||||
if value_type == 'int':
|
||||
if value is None or value == "":
|
||||
if value_type == "int":
|
||||
return 0
|
||||
elif value_type == 'float':
|
||||
elif value_type == "float":
|
||||
return 0.0
|
||||
elif value_type == 'bool':
|
||||
elif value_type == "bool":
|
||||
return False
|
||||
elif value_type == 'json':
|
||||
elif value_type == "json":
|
||||
return {}
|
||||
else:
|
||||
return ''
|
||||
|
||||
return ""
|
||||
|
||||
try:
|
||||
if value_type == 'int':
|
||||
if value_type == "int":
|
||||
return int(value)
|
||||
elif value_type == 'float':
|
||||
elif value_type == "float":
|
||||
return float(value)
|
||||
elif value_type == 'bool':
|
||||
return bool(value) or value == '1' or value == 'true' or value == 'True'
|
||||
elif value_type == 'json':
|
||||
elif value_type == "bool":
|
||||
return bool(value) or value == "1" or value == "true" or value == "True"
|
||||
elif value_type == "json":
|
||||
import json
|
||||
|
||||
try:
|
||||
return json.loads(value) if value else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return {}
|
||||
else:
|
||||
return str(value) if value is not None else ''
|
||||
return str(value) if value is not None else ""
|
||||
except (ValueError, TypeError) as e:
|
||||
logger.warning(f"Config value type conversion failed: value={value}, type={value_type}, error={str(e)}")
|
||||
# Returns default value if conversion fails
|
||||
if value_type == 'int':
|
||||
if value_type == "int":
|
||||
return 0
|
||||
elif value_type == 'float':
|
||||
elif value_type == "float":
|
||||
return 0.0
|
||||
elif value_type == 'bool':
|
||||
elif value_type == "bool":
|
||||
return False
|
||||
elif value_type == 'json':
|
||||
elif value_type == "json":
|
||||
return {}
|
||||
else:
|
||||
return str(value) if value is not None else ''
|
||||
return str(value) if value is not None else ""
|
||||
|
||||
|
||||
def get_internal_api_key() -> Optional[str]:
|
||||
"""
|
||||
Get the internal API key (preferably read from environment variables)
|
||||
|
||||
|
||||
Returns:
|
||||
Internal API key, returns None if not configured
|
||||
"""
|
||||
try:
|
||||
env_val = os.getenv('INTERNAL_API_KEY', '').strip()
|
||||
env_val = os.getenv("INTERNAL_API_KEY", "").strip()
|
||||
if env_val:
|
||||
return env_val
|
||||
|
||||
config = load_addon_config()
|
||||
api_key = config.get('internal_api', {}).get('key')
|
||||
|
||||
api_key = config.get("internal_api", {}).get("key")
|
||||
|
||||
if api_key:
|
||||
logger.debug(f"Loaded INTERNAL_API_KEY from env-config shape, length: {len(api_key)}")
|
||||
else:
|
||||
logger.warning("Missing INTERNAL_API_KEY (env).")
|
||||
|
||||
|
||||
return api_key
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load internal API key: {str(e)}")
|
||||
@@ -233,4 +219,3 @@ def clear_config_cache():
|
||||
global _config_cache
|
||||
_config_cache = None
|
||||
logger.debug("Addon config cache cleared")
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ Provides unified interface for PostgreSQL database operations.
|
||||
|
||||
Usage:
|
||||
from app.utils.db import get_db_connection
|
||||
|
||||
|
||||
with get_db_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM users WHERE id = %s", (user_id,))
|
||||
@@ -18,16 +18,22 @@ Configuration:
|
||||
|
||||
# Re-export from PostgreSQL module
|
||||
from app.utils.db_postgres import (
|
||||
get_pg_connection as get_db_connection,
|
||||
get_pg_connection_sync as get_db_connection_sync,
|
||||
is_postgres_available,
|
||||
close_pool as close_db,
|
||||
)
|
||||
from app.utils.db_postgres import (
|
||||
get_pg_connection as get_db_connection,
|
||||
)
|
||||
from app.utils.db_postgres import (
|
||||
get_pg_connection_sync as get_db_connection_sync,
|
||||
)
|
||||
from app.utils.db_postgres import (
|
||||
is_postgres_available,
|
||||
)
|
||||
|
||||
|
||||
def get_db_type() -> str:
|
||||
"""Get database type (always postgresql)"""
|
||||
return 'postgresql'
|
||||
return "postgresql"
|
||||
|
||||
|
||||
def is_postgres() -> bool:
|
||||
@@ -42,6 +48,7 @@ def init_database():
|
||||
"""
|
||||
if is_postgres_available():
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
logger.info("PostgreSQL connection verified")
|
||||
else:
|
||||
@@ -55,11 +62,11 @@ def close_db_connection():
|
||||
|
||||
|
||||
__all__ = [
|
||||
'get_db_connection',
|
||||
'get_db_connection_sync',
|
||||
'close_db_connection',
|
||||
'init_database',
|
||||
'close_db',
|
||||
'get_db_type',
|
||||
'is_postgres',
|
||||
"get_db_connection",
|
||||
"get_db_connection_sync",
|
||||
"close_db_connection",
|
||||
"init_database",
|
||||
"close_db",
|
||||
"get_db_type",
|
||||
"is_postgres",
|
||||
]
|
||||
|
||||
@@ -4,19 +4,21 @@ PostgreSQL Database Connection Utility
|
||||
Supports multi-user mode with connection pooling.
|
||||
Provides placeholder conversion for backward compatibility with legacy code.
|
||||
"""
|
||||
|
||||
import os
|
||||
import threading
|
||||
from typing import Optional, Any, List, Dict
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Try to import psycopg2
|
||||
try:
|
||||
import psycopg2
|
||||
from psycopg2 import pool
|
||||
from psycopg2.extras import RealDictCursor
|
||||
|
||||
HAS_PSYCOPG2 = True
|
||||
except ImportError:
|
||||
HAS_PSYCOPG2 = False
|
||||
@@ -29,7 +31,7 @@ _pool_lock = threading.Lock()
|
||||
|
||||
def _get_database_url() -> str:
|
||||
"""Get database connection URL from environment"""
|
||||
return os.getenv('DATABASE_URL', '').strip()
|
||||
return os.getenv("DATABASE_URL", "").strip()
|
||||
|
||||
|
||||
def _parse_database_url(url: str) -> Dict[str, Any]:
|
||||
@@ -38,131 +40,133 @@ def _parse_database_url(url: str) -> Dict[str, Any]:
|
||||
"""
|
||||
if not url:
|
||||
return {}
|
||||
|
||||
|
||||
# Remove protocol prefix
|
||||
if url.startswith('postgresql://'):
|
||||
if url.startswith("postgresql://"):
|
||||
url = url[13:]
|
||||
elif url.startswith('postgres://'):
|
||||
elif url.startswith("postgres://"):
|
||||
url = url[11:]
|
||||
else:
|
||||
return {}
|
||||
|
||||
|
||||
result = {}
|
||||
|
||||
|
||||
# Split user:password@host:port/dbname
|
||||
if '@' in url:
|
||||
auth, hostpart = url.rsplit('@', 1)
|
||||
if ':' in auth:
|
||||
result['user'], result['password'] = auth.split(':', 1)
|
||||
if "@" in url:
|
||||
auth, hostpart = url.rsplit("@", 1)
|
||||
if ":" in auth:
|
||||
result["user"], result["password"] = auth.split(":", 1)
|
||||
else:
|
||||
result['user'] = auth
|
||||
result["user"] = auth
|
||||
else:
|
||||
hostpart = url
|
||||
|
||||
|
||||
# Split host:port/dbname
|
||||
if '/' in hostpart:
|
||||
hostport, result['dbname'] = hostpart.split('/', 1)
|
||||
if "/" in hostpart:
|
||||
hostport, result["dbname"] = hostpart.split("/", 1)
|
||||
else:
|
||||
hostport = hostpart
|
||||
|
||||
if ':' in hostport:
|
||||
result['host'], port_str = hostport.split(':', 1)
|
||||
result['port'] = int(port_str)
|
||||
|
||||
if ":" in hostport:
|
||||
result["host"], port_str = hostport.split(":", 1)
|
||||
result["port"] = int(port_str)
|
||||
else:
|
||||
result['host'] = hostport
|
||||
result['port'] = 5432
|
||||
|
||||
result["host"] = hostport
|
||||
result["port"] = 5432
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _get_connection_pool():
|
||||
"""Get or create connection pool"""
|
||||
global _connection_pool
|
||||
|
||||
|
||||
if _connection_pool is not None:
|
||||
return _connection_pool
|
||||
|
||||
|
||||
with _pool_lock:
|
||||
if _connection_pool is not None:
|
||||
return _connection_pool
|
||||
|
||||
|
||||
if not HAS_PSYCOPG2:
|
||||
raise RuntimeError("psycopg2 is not installed. Cannot use PostgreSQL.")
|
||||
|
||||
|
||||
db_url = _get_database_url()
|
||||
if not db_url:
|
||||
raise RuntimeError("DATABASE_URL environment variable is not set.")
|
||||
|
||||
|
||||
params = _parse_database_url(db_url)
|
||||
if not params:
|
||||
raise RuntimeError(f"Invalid DATABASE_URL format: {db_url}")
|
||||
|
||||
|
||||
try:
|
||||
_connection_pool = pool.ThreadedConnectionPool(
|
||||
minconn=2,
|
||||
maxconn=20,
|
||||
host=params.get('host', 'localhost'),
|
||||
port=params.get('port', 5432),
|
||||
user=params.get('user', 'quantdinger'),
|
||||
password=params.get('password', ''),
|
||||
dbname=params.get('dbname', 'quantdinger'),
|
||||
host=params.get("host", "localhost"),
|
||||
port=params.get("port", 5432),
|
||||
user=params.get("user", "quantdinger"),
|
||||
password=params.get("password", ""),
|
||||
dbname=params.get("dbname", "quantdinger"),
|
||||
connect_timeout=10,
|
||||
)
|
||||
logger.info(f"PostgreSQL connection pool created: {params.get('host')}:{params.get('port')}/{params.get('dbname')}")
|
||||
logger.info(
|
||||
f"PostgreSQL connection pool created: {params.get('host')}:{params.get('port')}/{params.get('dbname')}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create PostgreSQL connection pool: {e}")
|
||||
raise
|
||||
|
||||
|
||||
return _connection_pool
|
||||
|
||||
|
||||
class PostgresCursor:
|
||||
"""PostgreSQL cursor wrapper with placeholder conversion for backward compatibility"""
|
||||
|
||||
|
||||
def __init__(self, cursor):
|
||||
self._cursor = cursor
|
||||
self._last_insert_id = None
|
||||
|
||||
|
||||
def _convert_placeholders(self, query: str) -> str:
|
||||
"""
|
||||
Convert ? placeholders to PostgreSQL %s for backward compatibility.
|
||||
Also handle some SQL syntax differences.
|
||||
"""
|
||||
# Replace ? -> %s
|
||||
query = query.replace('?', '%s')
|
||||
|
||||
query = query.replace("?", "%s")
|
||||
|
||||
# INSERT OR IGNORE -> PostgreSQL: INSERT ... ON CONFLICT DO NOTHING
|
||||
query = query.replace('INSERT OR IGNORE', 'INSERT')
|
||||
|
||||
query = query.replace("INSERT OR IGNORE", "INSERT")
|
||||
|
||||
return query
|
||||
|
||||
|
||||
def execute(self, query: str, args: Any = None):
|
||||
"""Execute SQL statement"""
|
||||
query = self._convert_placeholders(query)
|
||||
|
||||
|
||||
# Check if this is an INSERT and add RETURNING id if not present
|
||||
is_insert = query.strip().upper().startswith('INSERT')
|
||||
if is_insert and 'RETURNING' not in query.upper():
|
||||
query = query.rstrip(';').rstrip() + ' RETURNING id'
|
||||
|
||||
is_insert = query.strip().upper().startswith("INSERT")
|
||||
if is_insert and "RETURNING" not in query.upper():
|
||||
query = query.rstrip(";").rstrip() + " RETURNING id"
|
||||
|
||||
if args:
|
||||
if not isinstance(args, (tuple, list)):
|
||||
args = (args,)
|
||||
result = self._cursor.execute(query, args)
|
||||
else:
|
||||
result = self._cursor.execute(query)
|
||||
|
||||
|
||||
# Capture last insert id for INSERT statements
|
||||
if is_insert:
|
||||
try:
|
||||
row = self._cursor.fetchone()
|
||||
if row and 'id' in row:
|
||||
self._last_insert_id = row['id']
|
||||
if row and "id" in row:
|
||||
self._last_insert_id = row["id"]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def fetchone(self) -> Optional[Dict[str, Any]]:
|
||||
"""Fetch single row"""
|
||||
row = self._cursor.fetchone()
|
||||
@@ -170,7 +174,7 @@ class PostgresCursor:
|
||||
return None
|
||||
# RealDictCursor already returns a dict, so return as-is
|
||||
return row if isinstance(row, dict) else dict(row) if row else None
|
||||
|
||||
|
||||
def fetchall(self) -> List[Dict[str, Any]]:
|
||||
"""Fetch all rows"""
|
||||
rows = self._cursor.fetchall()
|
||||
@@ -178,16 +182,16 @@ class PostgresCursor:
|
||||
return []
|
||||
# RealDictCursor already returns dicts, so return as-is
|
||||
return [row if isinstance(row, dict) else dict(row) for row in rows]
|
||||
|
||||
|
||||
def close(self):
|
||||
"""Close cursor"""
|
||||
self._cursor.close()
|
||||
|
||||
|
||||
@property
|
||||
def lastrowid(self) -> Optional[int]:
|
||||
"""Get last inserted row ID"""
|
||||
return self._last_insert_id
|
||||
|
||||
|
||||
@property
|
||||
def rowcount(self) -> int:
|
||||
"""Get affected row count"""
|
||||
@@ -196,23 +200,23 @@ class PostgresCursor:
|
||||
|
||||
class PostgresConnection:
|
||||
"""PostgreSQL connection wrapper"""
|
||||
|
||||
|
||||
def __init__(self, conn):
|
||||
self._conn = conn
|
||||
self._pool = _get_connection_pool()
|
||||
|
||||
|
||||
def cursor(self) -> PostgresCursor:
|
||||
"""Create cursor"""
|
||||
return PostgresCursor(self._conn.cursor(cursor_factory=RealDictCursor))
|
||||
|
||||
|
||||
def commit(self):
|
||||
"""Commit transaction"""
|
||||
self._conn.commit()
|
||||
|
||||
|
||||
def rollback(self):
|
||||
"""Rollback transaction"""
|
||||
self._conn.rollback()
|
||||
|
||||
|
||||
def close(self):
|
||||
"""Return connection to pool"""
|
||||
if self._pool and self._conn:
|
||||
@@ -280,7 +284,7 @@ def execute_sql(sql: str, params: tuple = None) -> List[Dict[str, Any]]:
|
||||
with get_pg_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(sql, params)
|
||||
if sql.strip().upper().startswith('SELECT'):
|
||||
if sql.strip().upper().startswith("SELECT"):
|
||||
return cursor.fetchall()
|
||||
conn.commit()
|
||||
return []
|
||||
@@ -290,11 +294,11 @@ def is_postgres_available() -> bool:
|
||||
"""Check if PostgreSQL is available"""
|
||||
if not HAS_PSYCOPG2:
|
||||
return False
|
||||
|
||||
|
||||
db_url = _get_database_url()
|
||||
if not db_url:
|
||||
return False
|
||||
|
||||
|
||||
try:
|
||||
with get_pg_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
@@ -1,24 +1,23 @@
|
||||
"""
|
||||
HTTP tool module
|
||||
"""
|
||||
|
||||
import requests
|
||||
from requests.adapters import HTTPAdapter
|
||||
from urllib3.util.retry import Retry
|
||||
|
||||
|
||||
def get_retry_session(
|
||||
retries: int = 3,
|
||||
backoff_factor: float = 0.5,
|
||||
status_forcelist: tuple = (500, 502, 503, 504)
|
||||
retries: int = 3, backoff_factor: float = 0.5, status_forcelist: tuple = (500, 502, 503, 504)
|
||||
) -> requests.Session:
|
||||
"""
|
||||
Get HTTP Session with retry mechanism
|
||||
|
||||
|
||||
Args:
|
||||
retries: number of retries
|
||||
backoff_factor: retry interval factor
|
||||
status_forcelist: HTTP status codes that need to be retried
|
||||
|
||||
|
||||
Returns:
|
||||
Configured Session instance
|
||||
"""
|
||||
@@ -31,11 +30,10 @@ def get_retry_session(
|
||||
status_forcelist=status_forcelist,
|
||||
)
|
||||
adapter = HTTPAdapter(max_retries=retry)
|
||||
session.mount('http://', adapter)
|
||||
session.mount('https://', adapter)
|
||||
session.mount("http://", adapter)
|
||||
session.mount("https://", adapter)
|
||||
return session
|
||||
|
||||
|
||||
# Global shared session
|
||||
global_session = get_retry_session()
|
||||
|
||||
|
||||
@@ -9,7 +9,6 @@ from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
|
||||
SUPPORTED_LANGS = {
|
||||
"en-US",
|
||||
"zh-CN",
|
||||
@@ -82,5 +81,3 @@ def detect_request_language(flask_request, body: Optional[dict] = None, default:
|
||||
return lang
|
||||
|
||||
return default
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""
|
||||
Logging utilities (local-only friendly).
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from logging.handlers import RotatingFileHandler
|
||||
@@ -8,35 +9,32 @@ from logging.handlers import RotatingFileHandler
|
||||
|
||||
def setup_logger():
|
||||
"""Configure global logs"""
|
||||
log_level = os.getenv('LOG_LEVEL', 'INFO')
|
||||
log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, log_level.upper()),
|
||||
format=log_format
|
||||
)
|
||||
|
||||
log_level = os.getenv("LOG_LEVEL", "INFO")
|
||||
log_format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
|
||||
logging.basicConfig(level=getattr(logging, log_level.upper()), format=log_format)
|
||||
|
||||
# Filter werkzeug's INFO level logs (reduce noise)
|
||||
# Only keep WARNING and above levels
|
||||
werkzeug_logger = logging.getLogger('werkzeug')
|
||||
werkzeug_logger = logging.getLogger("werkzeug")
|
||||
werkzeug_logger.setLevel(logging.WARNING)
|
||||
|
||||
|
||||
# Filter INFO level logs for kline routes (reduce noise)
|
||||
# Only keep WARNING and above levels
|
||||
kline_logger = logging.getLogger('app.routes.kline')
|
||||
kline_logger = logging.getLogger("app.routes.kline")
|
||||
kline_logger.setLevel(logging.WARNING)
|
||||
|
||||
|
||||
# Create log directory
|
||||
log_dir = 'logs'
|
||||
log_dir = "logs"
|
||||
if not os.path.exists(log_dir):
|
||||
os.makedirs(log_dir)
|
||||
|
||||
|
||||
# Add file handler
|
||||
file_handler = RotatingFileHandler(
|
||||
os.path.join(log_dir, 'app.log'),
|
||||
maxBytes=10*1024*1024, # 10MB
|
||||
os.path.join(log_dir, "app.log"),
|
||||
maxBytes=10 * 1024 * 1024, # 10MB
|
||||
backupCount=5,
|
||||
encoding='utf-8'
|
||||
encoding="utf-8",
|
||||
)
|
||||
file_handler.setFormatter(logging.Formatter(log_format))
|
||||
logging.getLogger().addHandler(file_handler)
|
||||
@@ -45,12 +43,11 @@ def setup_logger():
|
||||
def get_logger(name: str) -> logging.Logger:
|
||||
"""
|
||||
Get the logger with the specified name
|
||||
|
||||
|
||||
Args:
|
||||
name: logger name
|
||||
|
||||
|
||||
Returns:
|
||||
Logger instance
|
||||
"""
|
||||
return logging.getLogger(name)
|
||||
|
||||
|
||||
@@ -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