123 lines
4.1 KiB
Python
123 lines
4.1 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
脚本输出捕获工具
|
|||
|
|
用于捕获脚本执行过程中的标准输出和错误输出，并通过WebSocket推送给订阅的客户端
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import sys
|
|||
|
|
from io import StringIO
|
|||
|
|
from contextlib import contextmanager
|
|||
|
|
from .script_websocket_service import get_websocket_manager
|
|||
|
|
from utils.logger import get_logger
|
|||
|
|
|
|||
|
|
logger = get_logger("services.script_output_capture")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ScriptOutputCapture:
|
|||
|
|
"""脚本输出捕获类"""
|
|||
|
|
|
|||
|
|
def __init__(self, script_id: str, level: str = "INFO", original_stream=None):
|
|||
|
|
self.script_id = script_id
|
|||
|
|
self.level = level
|
|||
|
|
self.websocket_manager = get_websocket_manager()
|
|||
|
|
self.buffer = StringIO()
|
|||
|
|
self.original_stream = original_stream # 保留原始输出流
|
|||
|
|
|
|||
|
|
def write(self, text):
|
|||
|
|
"""写入文本，同时推送到WebSocket和原始输出流"""
|
|||
|
|
# 1. 写入原始输出流（如控制台）
|
|||
|
|
if self.original_stream:
|
|||
|
|
try:
|
|||
|
|
self.original_stream.write(text)
|
|||
|
|
self.original_stream.flush()
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.debug(f"写入原始输出流失败: {e}")
|
|||
|
|
|
|||
|
|
# 2. 推送到WebSocket（仅对非空内容）
|
|||
|
|
if text and text.strip():
|
|||
|
|
try:
|
|||
|
|
# 尝试在当前事件循环中创建任务
|
|||
|
|
asyncio.create_task(self._push_output(text.strip()))
|
|||
|
|
except RuntimeError:
|
|||
|
|
# 如果没有运行中的事件循环，记录日志但不报错
|
|||
|
|
logger.debug(f"无法推送输出到WebSocket，没有运行中的事件循环: [{self.script_id}] {text.strip()}")
|
|||
|
|
|
|||
|
|
# 3. 写入缓冲区保留原有行为
|
|||
|
|
self.buffer.write(text)
|
|||
|
|
|
|||
|
|
def flush(self):
|
|||
|
|
"""刷新缓冲区"""
|
|||
|
|
self.buffer.flush()
|
|||
|
|
if self.original_stream:
|
|||
|
|
try:
|
|||
|
|
self.original_stream.flush()
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
async def _push_output(self, message: str):
|
|||
|
|
"""推送输出到WebSocket"""
|
|||
|
|
try:
|
|||
|
|
await self.websocket_manager.broadcast_script_log(
|
|||
|
|
self.script_id, self.level, message
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.debug(f"推送脚本输出失败 [{self.script_id}]: {e}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
@contextmanager
|
|||
|
|
def capture_script_output(script_id: str):
|
|||
|
|
"""
|
|||
|
|
上下文管理器，用于捕获脚本执行过程中的标准输出和错误输出
|
|||
|
|
|
|||
|
|
使用示例:
|
|||
|
|
with capture_script_output("script_123"):
|
|||
|
|
print("这条消息会被捕获并推送到WebSocket，同时显示在控制台")
|
|||
|
|
# 执行脚本代码
|
|||
|
|
"""
|
|||
|
|
# 备份原始输出流
|
|||
|
|
original_stdout = sys.stdout
|
|||
|
|
original_stderr = sys.stderr
|
|||
|
|
|
|||
|
|
# 创建捕获器，传入原始输出流以保持双重输出
|
|||
|
|
stdout_capture = ScriptOutputCapture(script_id, "INFO", original_stdout)
|
|||
|
|
stderr_capture = ScriptOutputCapture(script_id, "ERROR", original_stderr)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 替换标准输出流
|
|||
|
|
sys.stdout = stdout_capture
|
|||
|
|
sys.stderr = stderr_capture
|
|||
|
|
|
|||
|
|
yield {
|
|||
|
|
'stdout_capture': stdout_capture,
|
|||
|
|
'stderr_capture': stderr_capture
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
finally:
|
|||
|
|
# 恢复原始输出流
|
|||
|
|
sys.stdout = original_stdout
|
|||
|
|
sys.stderr = original_stderr
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def push_script_error(script_id: str, error_message: str):
|
|||
|
|
"""推送脚本错误信息到WebSocket"""
|
|||
|
|
try:
|
|||
|
|
websocket_manager = get_websocket_manager()
|
|||
|
|
await websocket_manager.broadcast_script_log(
|
|||
|
|
script_id, "ERROR", error_message
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.debug(f"推送脚本错误信息失败 [{script_id}]: {e}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def push_script_info(script_id: str, info_message: str):
|
|||
|
|
"""推送脚本信息到WebSocket"""
|
|||
|
|
try:
|
|||
|
|
websocket_manager = get_websocket_manager()
|
|||
|
|
await websocket_manager.broadcast_script_log(
|
|||
|
|
script_id, "INFO", info_message
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.debug(f"推送脚本信息失败 [{script_id}]: {e}")
|