#!/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 - 订阅脚本日志") print(" unsub - 取消订阅脚本日志") 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)