294 lines
12 KiB
Python
294 lines
12 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
|
||
|
|
"""
|
||
|
|
脚本与任务系统集成模块
|
||
|
|
提供动态路由处理和VWED任务中脚本模块引用运行
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
from typing import Dict, Any, Optional, Callable, List
|
||
|
|
from fastapi import Request, HTTPException
|
||
|
|
from fastapi.routing import APIRoute
|
||
|
|
from services.script_registry_service import get_global_registry
|
||
|
|
from services.script_engine_service import get_script_engine
|
||
|
|
from services.task_service import get_task_service
|
||
|
|
from utils.logger import get_logger
|
||
|
|
|
||
|
|
logger = get_logger("services.script_task_integration")
|
||
|
|
|
||
|
|
|
||
|
|
class DynamicRouteHandler:
|
||
|
|
"""动态路由处理器 - 处理脚本注册的API接口"""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self.registry = get_global_registry()
|
||
|
|
self.script_engine = get_script_engine()
|
||
|
|
|
||
|
|
async def handle_dynamic_route(self, request: Request) -> Dict[str, Any]:
|
||
|
|
"""处理动态路由请求"""
|
||
|
|
try:
|
||
|
|
# 获取请求信息
|
||
|
|
path = request.url.path
|
||
|
|
method = request.method
|
||
|
|
|
||
|
|
# 查找对应的API路由
|
||
|
|
route_key = f"{method}:{path}"
|
||
|
|
route_info = self.registry.get_api_route(path, method)
|
||
|
|
|
||
|
|
if not route_info:
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=404,
|
||
|
|
detail=f"API路由未找到: {method} {path}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# 检查脚本是否还在运行
|
||
|
|
script_id = route_info["script_id"]
|
||
|
|
running_scripts = self.script_engine.get_running_scripts()
|
||
|
|
script_running = any(s["script_id"] == script_id for s in running_scripts)
|
||
|
|
|
||
|
|
if not script_running:
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=503,
|
||
|
|
detail=f"脚本 {script_id} 未运行，无法处理请求"
|
||
|
|
)
|
||
|
|
|
||
|
|
# 获取请求数据
|
||
|
|
request_data = await self._extract_request_data(request)
|
||
|
|
|
||
|
|
# 调用处理函数
|
||
|
|
start_time = asyncio.get_event_loop().time()
|
||
|
|
handler = route_info["handler"]
|
||
|
|
|
||
|
|
try:
|
||
|
|
# 检查是否为异步函数
|
||
|
|
if asyncio.iscoroutinefunction(handler):
|
||
|
|
result = await handler(request_data)
|
||
|
|
else:
|
||
|
|
result = handler(request_data)
|
||
|
|
|
||
|
|
# 更新调用统计
|
||
|
|
response_time_ms = int((asyncio.get_event_loop().time() - start_time) * 1000)
|
||
|
|
self.registry.update_api_call_stats(route_key, response_time_ms)
|
||
|
|
|
||
|
|
return {"success": True, "data": result}
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"API处理函数执行失败: {route_key} -> {e}", exc_info=True)
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=500,
|
||
|
|
detail=f"API处理失败: {str(e)}"
|
||
|
|
)
|
||
|
|
|
||
|
|
except HTTPException:
|
||
|
|
raise
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"动态路由处理失败: {e}", exc_info=True)
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=500,
|
||
|
|
detail=f"路由处理失败: {str(e)}"
|
||
|
|
)
|
||
|
|
|
||
|
|
async def _extract_request_data(self, request: Request) -> Dict[str, Any]:
|
||
|
|
"""提取请求数据"""
|
||
|
|
try:
|
||
|
|
# 获取查询参数
|
||
|
|
query_params = dict(request.query_params)
|
||
|
|
|
||
|
|
# 获取路径参数
|
||
|
|
path_params = dict(request.path_params)
|
||
|
|
|
||
|
|
# 获取请求体数据
|
||
|
|
body_data = {}
|
||
|
|
if request.method in ["POST", "PUT", "PATCH"]:
|
||
|
|
content_type = request.headers.get("content-type", "")
|
||
|
|
if "application/json" in content_type:
|
||
|
|
body_data = await request.json()
|
||
|
|
elif "application/x-www-form-urlencoded" in content_type:
|
||
|
|
form_data = await request.form()
|
||
|
|
body_data = dict(form_data)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"method": request.method,
|
||
|
|
"path": request.url.path,
|
||
|
|
"query_params": query_params,
|
||
|
|
"path_params": path_params,
|
||
|
|
"body": body_data,
|
||
|
|
"headers": dict(request.headers),
|
||
|
|
"client": request.client.host if request.client else None
|
||
|
|
}
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"提取请求数据失败: {e}", exc_info=True)
|
||
|
|
return {}
|
||
|
|
|
||
|
|
|
||
|
|
class ScriptTaskIntegration:
|
||
|
|
"""脚本任务集成服务"""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self.registry = get_global_registry()
|
||
|
|
self.script_engine = get_script_engine()
|
||
|
|
self.task_service = get_task_service()
|
||
|
|
|
||
|
|
async def execute_script_function_from_task(self, function_name: str,
|
||
|
|
function_args: Any,
|
||
|
|
task_context: Dict[str, Any] = None) -> Dict[str, Any]:
|
||
|
|
"""从VWED任务中执行脚本函数"""
|
||
|
|
try:
|
||
|
|
# 查找函数信息
|
||
|
|
func_info = self.registry.get_function_info(function_name)
|
||
|
|
if not func_info:
|
||
|
|
return {
|
||
|
|
"success": False,
|
||
|
|
"error": f"函数 {function_name} 未注册"
|
||
|
|
}
|
||
|
|
|
||
|
|
script_id = func_info["script_id"]
|
||
|
|
|
||
|
|
# 检查脚本是否还在运行
|
||
|
|
running_scripts = self.script_engine.get_running_scripts()
|
||
|
|
script_running = any(s["script_id"] == script_id for s in running_scripts)
|
||
|
|
|
||
|
|
if not script_running:
|
||
|
|
return {
|
||
|
|
"success": False,
|
||
|
|
"error": f"脚本 {script_id} 未运行，无法执行函数"
|
||
|
|
}
|
||
|
|
|
||
|
|
# 执行函数
|
||
|
|
result = await self.script_engine.execute_script_function(
|
||
|
|
script_id=script_id,
|
||
|
|
function_name=function_name,
|
||
|
|
function_args=function_args
|
||
|
|
)
|
||
|
|
|
||
|
|
# 如果有任务上下文，记录执行结果
|
||
|
|
if task_context:
|
||
|
|
await self._record_task_execution(
|
||
|
|
task_context, function_name, function_args, result
|
||
|
|
)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"从任务执行脚本函数失败: {e}", exc_info=True)
|
||
|
|
return {
|
||
|
|
"success": False,
|
||
|
|
"error": f"执行脚本函数失败: {str(e)}"
|
||
|
|
}
|
||
|
|
|
||
|
|
async def get_available_functions(self, tags: list = None) -> List[Dict[str, Any]]:
|
||
|
|
"""获取可用的脚本函数列表（供任务编辑器使用）"""
|
||
|
|
try:
|
||
|
|
all_registrations = self.registry.get_all_registrations()
|
||
|
|
functions = all_registrations.get("registered_functions", {})
|
||
|
|
|
||
|
|
function_list = []
|
||
|
|
for func_name, func_info in functions.items():
|
||
|
|
# 检查脚本是否运行
|
||
|
|
script_id = func_info["script_id"]
|
||
|
|
running_scripts = self.script_engine.get_running_scripts()
|
||
|
|
is_running = any(s["script_id"] == script_id for s in running_scripts)
|
||
|
|
|
||
|
|
# 标签筛选
|
||
|
|
if tags:
|
||
|
|
func_tags = func_info.get("tags", [])
|
||
|
|
if not any(tag in func_tags for tag in tags):
|
||
|
|
continue
|
||
|
|
|
||
|
|
function_info = {
|
||
|
|
"function_name": func_name,
|
||
|
|
"script_id": script_id,
|
||
|
|
"description": func_info.get("description", ""),
|
||
|
|
"parameters": func_info.get("parameters", []),
|
||
|
|
"return_schema": func_info.get("return_schema", {}),
|
||
|
|
"is_async": func_info.get("is_async", False),
|
||
|
|
"tags": func_info.get("tags", []),
|
||
|
|
"is_running": is_running,
|
||
|
|
"stats": {
|
||
|
|
"call_count": func_info.get("call_count", 0),
|
||
|
|
"success_count": func_info.get("success_count", 0),
|
||
|
|
"error_count": func_info.get("error_count", 0),
|
||
|
|
"average_execution_time_ms": func_info.get("average_execution_time_ms")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
function_list.append(function_info)
|
||
|
|
|
||
|
|
return function_list
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"获取可用函数列表失败: {e}", exc_info=True)
|
||
|
|
return []
|
||
|
|
|
||
|
|
async def register_script_block_handler(self):
|
||
|
|
"""注册脚本块处理器到现有任务系统"""
|
||
|
|
try:
|
||
|
|
# 这里需要根据现有的任务系统架构进行集成
|
||
|
|
# 扩展现有的 ScriptHandler 来支持动态函数调用
|
||
|
|
|
||
|
|
from services.execution.handlers.script import ScriptBlockHandler
|
||
|
|
|
||
|
|
# 增强脚本处理器
|
||
|
|
original_execute = ScriptBlockHandler.execute
|
||
|
|
|
||
|
|
async def enhanced_execute(self, block_config: Dict[str, Any], context) -> Dict[str, Any]:
|
||
|
|
"""增强的脚本执行方法"""
|
||
|
|
try:
|
||
|
|
# 检查是否为函数调用模式
|
||
|
|
if block_config.get("execution_mode") == "function_call":
|
||
|
|
function_name = block_config.get("function_name")
|
||
|
|
function_args = block_config.get("function_args", {})
|
||
|
|
|
||
|
|
if not function_name:
|
||
|
|
return {"success": False, "error": "未指定函数名"}
|
||
|
|
|
||
|
|
# 使用集成服务执行函数
|
||
|
|
integration = ScriptTaskIntegration()
|
||
|
|
result = await integration.execute_script_function_from_task(
|
||
|
|
function_name=function_name,
|
||
|
|
function_args=function_args,
|
||
|
|
task_context=context.to_dict() if hasattr(context, 'to_dict') else {}
|
||
|
|
)
|
||
|
|
|
||
|
|
return result
|
||
|
|
else:
|
||
|
|
# 使用原始处理逻辑
|
||
|
|
return await original_execute(self, block_config, context)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"增强脚本处理器执行失败: {e}", exc_info=True)
|
||
|
|
return {"success": False, "error": f"脚本执行失败: {str(e)}"}
|
||
|
|
|
||
|
|
# 替换原始方法
|
||
|
|
ScriptBlockHandler.execute = enhanced_execute
|
||
|
|
logger.info("脚本块处理器已增强，支持动态函数调用")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"注册脚本块处理器失败: {e}", exc_info=True)
|
||
|
|
|
||
|
|
async def _record_task_execution(self, task_context: Dict[str, Any],
|
||
|
|
function_name: str, function_args: Any,
|
||
|
|
execution_result: Dict[str, Any]):
|
||
|
|
"""记录任务执行信息"""
|
||
|
|
try:
|
||
|
|
# 这里可以记录到任务执行日志或其他监控系统
|
||
|
|
logger.info(f"任务脚本函数执行: {function_name} -> {execution_result.get('success', False)}")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"记录任务执行信息失败: {e}", exc_info=True)
|
||
|
|
|
||
|
|
|
||
|
|
# 全局集成服务实例
|
||
|
|
_task_integration = ScriptTaskIntegration()
|
||
|
|
_route_handler = DynamicRouteHandler()
|
||
|
|
|
||
|
|
|
||
|
|
def get_task_integration() -> ScriptTaskIntegration:
|
||
|
|
"""获取脚本任务集成服务实例"""
|
||
|
|
return _task_integration
|
||
|
|
|
||
|
|
|
||
|
|
def get_dynamic_route_handler() -> DynamicRouteHandler:
|
||
|
|
"""获取动态路由处理器实例"""
|
||
|
|
return _route_handler
|