320 lines
11 KiB
Python
320 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
WebSocket 脚本日志订阅测试
|
||
用于监控 VWED 脚本日志推送
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
from datetime import datetime
|
||
import signal
|
||
import sys
|
||
|
||
try:
|
||
import websockets
|
||
except ImportError:
|
||
print("❌ 缺少 websockets 依赖，请运行: pip install websockets")
|
||
sys.exit(1)
|
||
|
||
|
||
class VWEDWebSocketLogMonitor:
|
||
"""VWED WebSocket 日志监控器"""
|
||
|
||
def __init__(self, host="localhost", port=8000, client_id="monitor_test", client_type="monitor"):
|
||
self.host = host
|
||
self.port = port
|
||
self.client_id = client_id
|
||
self.client_type = client_type
|
||
self.websocket = None
|
||
self.running = False
|
||
self.subscribed_scripts = set()
|
||
|
||
async def connect(self):
|
||
"""连接到 WebSocket 服务器"""
|
||
uri = f"ws://{self.host}:{self.port}/api/script/ws/script-logs?client_id={self.client_id}&client_type={self.client_type}"
|
||
|
||
try:
|
||
print(f"正在连接到 WebSocket 服务器: {uri}")
|
||
self.websocket = await websockets.connect(uri)
|
||
self.running = True
|
||
print(f"✅ WebSocket 连接成功建立")
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ WebSocket 连接失败: {e}")
|
||
return False
|
||
|
||
async def disconnect(self):
|
||
"""断开连接"""
|
||
self.running = False
|
||
if self.websocket:
|
||
await self.websocket.close()
|
||
print("🔌 WebSocket 连接已断开")
|
||
|
||
async def subscribe_script(self, script_id):
|
||
"""订阅脚本日志"""
|
||
if not self.websocket:
|
||
print("WebSocket 未连接")
|
||
return False
|
||
|
||
try:
|
||
message = {
|
||
"type": "subscribe",
|
||
"script_id": script_id
|
||
}
|
||
await self.websocket.send(json.dumps(message))
|
||
self.subscribed_scripts.add(script_id)
|
||
print(f"📡 已发送订阅请求: {script_id}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"订阅脚本失败: {e}")
|
||
return False
|
||
|
||
async def unsubscribe_script(self, script_id):
|
||
"""取消订阅脚本日志"""
|
||
if not self.websocket:
|
||
print("WebSocket 未连接")
|
||
return False
|
||
|
||
try:
|
||
message = {
|
||
"type": "unsubscribe",
|
||
"script_id": script_id
|
||
}
|
||
await self.websocket.send(json.dumps(message))
|
||
self.subscribed_scripts.discard(script_id)
|
||
print(f" 已发送取消订阅请求: {script_id}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"取消订阅脚本失败: {e}")
|
||
return False
|
||
|
||
async def get_status(self):
|
||
"""获取连接状态"""
|
||
if not self.websocket:
|
||
print("WebSocket 未连接")
|
||
return False
|
||
|
||
try:
|
||
message = {"type": "get_status"}
|
||
await self.websocket.send(json.dumps(message))
|
||
print("已发送状态查询请求")
|
||
return True
|
||
except Exception as e:
|
||
print(f"获取状态失败: {e}")
|
||
return False
|
||
|
||
async def send_ping(self):
|
||
"""发送心跳"""
|
||
if not self.websocket:
|
||
return False
|
||
|
||
try:
|
||
message = {
|
||
"type": "ping",
|
||
"timestamp": datetime.now().isoformat()
|
||
}
|
||
await self.websocket.send(json.dumps(message))
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ 发送心跳失败: {e}")
|
||
return False
|
||
|
||
def format_message(self, message):
|
||
"""格式化消息显示"""
|
||
msg_type = message.get("type", "unknown")
|
||
timestamp = datetime.now().strftime("%H:%M:%S")
|
||
|
||
if msg_type == "welcome":
|
||
return f"🎉 [{timestamp}] 欢迎消息: {message.get('message')} (连接ID: {message.get('connection_id')})"
|
||
|
||
elif msg_type == "script_log":
|
||
script_id = message.get("script_id", "unknown")
|
||
level = message.get("level", "INFO")
|
||
log_msg = message.get("message", "")
|
||
|
||
# 根据日志级别选择emoji
|
||
level_emoji = {
|
||
"INFO": "ℹ️",
|
||
"WARNING": "⚠️",
|
||
"ERROR": "❌",
|
||
"DEBUG": "🐛"
|
||
}
|
||
emoji = level_emoji.get(level, "📝")
|
||
|
||
return f"{emoji} [{timestamp}] [{script_id}] {level}: {log_msg}"
|
||
|
||
elif msg_type == "script_status":
|
||
script_id = message.get("script_id", "unknown")
|
||
status = message.get("status", "unknown")
|
||
status_msg = message.get("message", "")
|
||
return f"🔄 [{timestamp}] [{script_id}] 状态变更: {status} - {status_msg}"
|
||
|
||
elif msg_type == "function_execution":
|
||
script_id = message.get("script_id", "unknown")
|
||
func_name = message.get("function_name", "unknown")
|
||
result = message.get("result", {})
|
||
return f"⚡ [{timestamp}] [{script_id}] 函数执行: {func_name} -> {json.dumps(result, ensure_ascii=False)}"
|
||
|
||
elif msg_type == "subscription_success":
|
||
return f"✅ [{timestamp}] 订阅成功: {message.get('message')}"
|
||
|
||
elif msg_type == "unsubscription_success":
|
||
return f"✅ [{timestamp}] 取消订阅成功: {message.get('message')}"
|
||
|
||
elif msg_type == "status":
|
||
data = message.get("data", {})
|
||
total_connections = data.get("total_connections", 0)
|
||
total_scripts = data.get("total_script_subscriptions", 0)
|
||
return f"📊 [{timestamp}] 连接状态: {total_connections} 个连接, {total_scripts} 个脚本订阅"
|
||
|
||
elif msg_type == "error":
|
||
return f"❌ [{timestamp}] 错误: {message.get('message')}"
|
||
|
||
elif msg_type == "pong":
|
||
return f"💓 [{timestamp}] 心跳响应"
|
||
|
||
else:
|
||
return f"❓ [{timestamp}] 未知消息类型 {msg_type}: {json.dumps(message, ensure_ascii=False)}"
|
||
|
||
async def listen_messages(self):
|
||
"""监听消息"""
|
||
if not self.websocket:
|
||
print("❌ WebSocket 未连接")
|
||
return
|
||
|
||
try:
|
||
while self.running:
|
||
try:
|
||
# 等待消息，设置超时避免无限阻塞
|
||
message = await asyncio.wait_for(
|
||
self.websocket.recv(),
|
||
timeout=1.0
|
||
)
|
||
|
||
try:
|
||
data = json.loads(message)
|
||
formatted_msg = self.format_message(data)
|
||
print(formatted_msg)
|
||
except json.JSONDecodeError:
|
||
print(f"❌ JSON 解析失败: {message}")
|
||
|
||
except asyncio.TimeoutError:
|
||
# 超时是正常的，继续循环
|
||
continue
|
||
except websockets.exceptions.ConnectionClosed:
|
||
print("🔌 WebSocket 连接已关闭")
|
||
break
|
||
except Exception as e:
|
||
print(f"❌ 接收消息异常: {e}")
|
||
break
|
||
|
||
except Exception as e:
|
||
print(f"❌ 监听消息异常: {e}")
|
||
|
||
async def interactive_monitor(self):
|
||
"""交互式监控"""
|
||
print("\n" + "="*50)
|
||
print("VWED WebSocket 脚本日志监控器")
|
||
print("="*50)
|
||
print("可用命令:")
|
||
print(" sub <script_id> - 订阅脚本日志")
|
||
print(" unsub <script_id> - 取消订阅脚本日志")
|
||
print(" status - 获取连接状态")
|
||
print(" ping - 发送心跳")
|
||
print(" list - 显示已订阅的脚本")
|
||
print(" quit - 退出监控")
|
||
print("="*50)
|
||
|
||
# 启动消息监听任务
|
||
listen_task = asyncio.create_task(self.listen_messages())
|
||
|
||
try:
|
||
while self.running:
|
||
try:
|
||
# 等待用户输入
|
||
user_input = await asyncio.get_event_loop().run_in_executor(
|
||
None, input, "请输入命令 (sub/unsub/status/ping/list/quit): "
|
||
)
|
||
|
||
command_parts = user_input.strip().split()
|
||
if not command_parts:
|
||
continue
|
||
|
||
command = command_parts[0].lower()
|
||
|
||
if command == "quit":
|
||
print("👋 退出监控...")
|
||
break
|
||
elif command == "sub" and len(command_parts) > 1:
|
||
script_id = command_parts[1]
|
||
await self.subscribe_script(script_id)
|
||
elif command == "unsub" and len(command_parts) > 1:
|
||
script_id = command_parts[1]
|
||
await self.unsubscribe_script(script_id)
|
||
elif command == "status":
|
||
await self.get_status()
|
||
elif command == "ping":
|
||
await self.send_ping()
|
||
elif command == "list":
|
||
if self.subscribed_scripts:
|
||
print(f"📋 已订阅的脚本: {', '.join(self.subscribed_scripts)}")
|
||
else:
|
||
print("📋 当前没有订阅任何脚本")
|
||
else:
|
||
print("❌ 未知命令或参数不足")
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n👋 收到退出信号...")
|
||
break
|
||
except Exception as e:
|
||
print(f"❌ 处理命令异常: {e}")
|
||
|
||
finally:
|
||
self.running = False
|
||
listen_task.cancel()
|
||
try:
|
||
await listen_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
|
||
async def main():
|
||
"""主函数"""
|
||
monitor = VWEDWebSocketLogMonitor()
|
||
|
||
# 设置信号处理
|
||
def signal_handler(signum, _):
|
||
print(f"\n👋 收到信号 {signum}，正在退出...")
|
||
monitor.running = False
|
||
|
||
signal.signal(signal.SIGINT, signal_handler)
|
||
signal.signal(signal.SIGTERM, signal_handler)
|
||
|
||
try:
|
||
# 连接 WebSocket
|
||
if await monitor.connect():
|
||
# 开始交互式监控
|
||
await monitor.interactive_monitor()
|
||
else:
|
||
print("❌ 无法连接到 WebSocket 服务器")
|
||
return 1
|
||
|
||
except Exception as e:
|
||
print(f"❌ 监控异常: {e}")
|
||
return 1
|
||
|
||
finally:
|
||
# 清理连接
|
||
await monitor.disconnect()
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
exit_code = asyncio.run(main())
|
||
sys.exit(exit_code)
|
||
except KeyboardInterrupt:
|
||
print("\n👋 程序已退出")
|
||
sys.exit(0) |