Baseline: pr1 HYC下载站 v2.3 before security/functional fixes
This commit is contained in:
HYC Fixer
@@ -0,0 +1,403 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
WebSocket处理器模块
|
||||
提供双向实时通信功能
|
||||
"""
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Dict, Set, Optional, Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class WebSocketClient:
|
||||
"""WebSocket客户端"""
|
||||
client_id: str
|
||||
connection: any # WebSocket连接对象
|
||||
topics: Set[str] = field(default_factory=set)
|
||||
connected_at: float = field(default_factory=time.time)
|
||||
last_activity: float = field(default_factory=time.time)
|
||||
metadata: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
class WebSocketManager:
|
||||
"""WebSocket连接管理器"""
|
||||
|
||||
# 预定义主题
|
||||
TOPIC_MONITOR_CPU = 'monitor:cpu'
|
||||
TOPIC_MONITOR_MEMORY = 'monitor:memory'
|
||||
TOPIC_MONITOR_DISK = 'monitor:disk'
|
||||
TOPIC_MONITOR_NETWORK = 'monitor:network'
|
||||
TOPIC_MONITOR_ALL = 'monitor:*'
|
||||
TOPIC_SYNC_PROGRESS = 'sync:progress'
|
||||
TOPIC_SYNC_STATUS = 'sync:status'
|
||||
TOPIC_SYNC_ALL = 'sync:*'
|
||||
TOPIC_DOWNLOAD_PROGRESS = 'download:progress'
|
||||
TOPIC_DOWNLOAD_STATUS = 'download:status'
|
||||
TOPIC_DOWNLOAD_ALL = 'download:*'
|
||||
TOPIC_SERVER_STATUS = 'server:status'
|
||||
TOPIC_ALL = '*'
|
||||
|
||||
def __init__(self, config: dict = None):
|
||||
self.config = config or {}
|
||||
self.clients: Dict[str, WebSocketClient] = {}
|
||||
self.lock = threading.Lock()
|
||||
self._running = False
|
||||
|
||||
# 消息队列（用于批量发送）
|
||||
self.message_queues: Dict[str, list] = {}
|
||||
|
||||
# 心跳间隔（秒）
|
||||
self.heartbeat_interval = 30
|
||||
|
||||
# 统计
|
||||
self.stats = {
|
||||
'total_connections': 0,
|
||||
'total_messages_sent': 0,
|
||||
'total_messages_received': 0
|
||||
}
|
||||
|
||||
def register_client(self, client_id: str, connection, metadata: dict = None) -> WebSocketClient:
|
||||
"""注册新客户端"""
|
||||
with self.lock:
|
||||
client = WebSocketClient(
|
||||
client_id=client_id,
|
||||
connection=connection,
|
||||
topics={self.TOPIC_ALL}, # 默认订阅所有
|
||||
metadata=metadata or {}
|
||||
)
|
||||
self.clients[client_id] = client
|
||||
self.stats['total_connections'] += 1
|
||||
|
||||
# 启动心跳
|
||||
self._start_heartbeat(client_id)
|
||||
|
||||
return client
|
||||
|
||||
def unregister_client(self, client_id: str):
|
||||
"""注销客户端"""
|
||||
with self.lock:
|
||||
client = self.clients.pop(client_id, None)
|
||||
if client:
|
||||
self.message_queues.pop(client_id, None)
|
||||
|
||||
def subscribe(self, client_id: str, *topics: str):
|
||||
"""订阅主题"""
|
||||
with self.lock:
|
||||
client = self.clients.get(client_id)
|
||||
if client:
|
||||
client.topics.update(topics)
|
||||
|
||||
def unsubscribe(self, client_id: str, *topics: str):
|
||||
"""取消订阅"""
|
||||
with self.lock:
|
||||
client = self.clients.get(client_id)
|
||||
if client:
|
||||
for topic in topics:
|
||||
client.topics.discard(topic)
|
||||
|
||||
def is_subscribed(self, client_id: str, topic: str) -> bool:
|
||||
"""检查客户端是否订阅了主题"""
|
||||
with self.lock:
|
||||
client = self.clients.get(client_id)
|
||||
if not client:
|
||||
return False
|
||||
|
||||
if self.TOPIC_ALL in client.topics:
|
||||
return True
|
||||
|
||||
# 检查精确匹配或通配符匹配
|
||||
if topic in client.topics:
|
||||
return True
|
||||
|
||||
# 检查通配符匹配
|
||||
topic_parts = topic.split(':')
|
||||
for subscribed_topic in client.topics:
|
||||
if subscribed_topic.endswith('*'):
|
||||
prefix = subscribed_topic.rstrip('*').rstrip(':')
|
||||
if topic.startswith(prefix):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def send_to_client(self, client_id: str, event_type: str, data: dict, callback: Callable = None):
|
||||
"""发送消息到指定客户端"""
|
||||
with self.lock:
|
||||
client = self.clients.get(client_id)
|
||||
if not client:
|
||||
return False
|
||||
|
||||
message = self._format_message(event_type, data)
|
||||
|
||||
try:
|
||||
if hasattr(client.connection, 'send'):
|
||||
client.connection.send(message)
|
||||
self.stats['total_messages_sent'] += 1
|
||||
client.last_activity = time.time()
|
||||
|
||||
if callback:
|
||||
callback(client_id, True)
|
||||
return True
|
||||
else:
|
||||
# 放入消息队列
|
||||
if client_id not in self.message_queues:
|
||||
self.message_queues[client_id] = []
|
||||
self.message_queues[client_id].append(message)
|
||||
|
||||
if callback:
|
||||
callback(client_id, True)
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"WebSocket发送失败 {client_id}: {e}")
|
||||
if callback:
|
||||
callback(client_id, False)
|
||||
return False
|
||||
|
||||
def broadcast(self, event_type: str, data: dict, topic: str = None):
|
||||
"""广播消息到所有客户端"""
|
||||
with self.lock:
|
||||
sent_count = 0
|
||||
failed_clients = []
|
||||
|
||||
for client_id, client in self.clients.items():
|
||||
# 检查是否匹配主题
|
||||
if topic and not self.is_subscribed(client_id, topic):
|
||||
continue
|
||||
|
||||
message = self._format_message(event_type, data)
|
||||
|
||||
try:
|
||||
if hasattr(client.connection, 'send'):
|
||||
client.connection.send(message)
|
||||
sent_count += 1
|
||||
else:
|
||||
if client_id not in self.message_queues:
|
||||
self.message_queues[client_id] = []
|
||||
self.message_queues[client_id].append(message)
|
||||
sent_count += 1
|
||||
|
||||
client.last_activity = time.time()
|
||||
|
||||
except Exception as e:
|
||||
print(f"广播到 {client_id} 失败: {e}")
|
||||
failed_clients.append(client_id)
|
||||
|
||||
self.stats['total_messages_sent'] += sent_count
|
||||
|
||||
# 清理失败的客户端
|
||||
for client_id in failed_clients:
|
||||
self.unregister_client(client_id)
|
||||
|
||||
return sent_count
|
||||
|
||||
def broadcast_to_topic(self, topic: str, event_type: str, data: dict):
|
||||
"""广播到订阅特定主题的客户端"""
|
||||
with self.lock:
|
||||
sent_count = 0
|
||||
|
||||
for client_id, client in self.clients.items():
|
||||
if self.is_subscribed(client_id, topic):
|
||||
message = self._format_message(event_type, data)
|
||||
|
||||
try:
|
||||
if hasattr(client.connection, 'send'):
|
||||
client.connection.send(message)
|
||||
sent_count += 1
|
||||
else:
|
||||
if client_id not in self.message_queues:
|
||||
self.message_queues[client_id] = []
|
||||
self.message_queues[client_id].append(message)
|
||||
sent_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f"发送到 {client_id} 失败: {e}")
|
||||
|
||||
return sent_count
|
||||
|
||||
def broadcast_monitor_update(self, stats: dict):
|
||||
"""广播监控更新"""
|
||||
# 提取关键指标
|
||||
cpu_percent = stats.get('cpu', {}).get('percent', 0)
|
||||
memory_percent = stats.get('memory', {}).get('percent', 0)
|
||||
disk_percent = stats.get('disk', {}).get('percent', 0)
|
||||
|
||||
# 按阈值过滤
|
||||
if cpu_percent > 0 or memory_percent > 0 or disk_percent > 0:
|
||||
self.broadcast('monitor:stats', stats, 'monitor:*')
|
||||
|
||||
def broadcast_sync_progress(self, task_id: str, progress: dict):
|
||||
"""广播同步进度"""
|
||||
self.broadcast('sync:progress', progress, 'sync:*')
|
||||
|
||||
def broadcast_download_progress(self, download_id: str, progress: dict):
|
||||
"""广播下载进度"""
|
||||
self.broadcast('download:progress', progress, 'download:*')
|
||||
|
||||
def get_client_count(self) -> int:
|
||||
"""获取客户端数量"""
|
||||
with self.lock:
|
||||
return len(self.clients)
|
||||
|
||||
def get_client_topics(self, client_id: str) -> Set[str]:
|
||||
"""获取客户端订阅的主题"""
|
||||
with self.lock:
|
||||
client = self.clients.get(client_id)
|
||||
return client.topics.copy() if client else set()
|
||||
|
||||
def get_all_topics(self) -> Set[str]:
|
||||
"""获取所有被订阅的主题"""
|
||||
with self.lock:
|
||||
topics = set()
|
||||
for client in self.clients.values():
|
||||
topics.update(client.topics)
|
||||
return topics
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
"""获取统计信息"""
|
||||
with self.lock:
|
||||
return {
|
||||
**self.stats,
|
||||
'connected_clients': len(self.clients),
|
||||
'active_topics': len(self.get_all_topics()),
|
||||
'queues_queued': sum(len(q) for q in self.message_queues.values())
|
||||
}
|
||||
|
||||
def _format_message(self, event_type: str, data: dict) -> str:
|
||||
"""格式化消息"""
|
||||
return json.dumps({
|
||||
'type': event_type,
|
||||
'timestamp': time.time(),
|
||||
'data': data
|
||||
}, ensure_ascii=False)
|
||||
|
||||
def _start_heartbeat(self, client_id: str):
|
||||
"""启动心跳"""
|
||||
def heartbeat():
|
||||
while client_id in self.clients:
|
||||
try:
|
||||
# 发送心跳
|
||||
self.send_to_client(
|
||||
client_id,
|
||||
'ping',
|
||||
{'timestamp': time.time()}
|
||||
)
|
||||
except Exception:
|
||||
break
|
||||
|
||||
time.sleep(self.heartbeat_interval)
|
||||
|
||||
thread = threading.Thread(target=heartbeat, daemon=True)
|
||||
thread.start()
|
||||
|
||||
def generate_client_id(self) -> str:
|
||||
"""生成客户端ID"""
|
||||
return f"ws_{uuid.uuid4().hex[:12]}"
|
||||
|
||||
def process_messages(self, client_id: str, messages: list):
|
||||
"""处理客户端消息"""
|
||||
for message in messages:
|
||||
self._handle_message(client_id, message)
|
||||
|
||||
def _handle_message(self, client_id: str, message: str):
|
||||
"""处理客户端消息"""
|
||||
self.stats['total_messages_received'] += 1
|
||||
|
||||
try:
|
||||
data = json.loads(message)
|
||||
|
||||
event_type = data.get('type')
|
||||
payload = data.get('data', {})
|
||||
|
||||
if event_type == 'subscribe':
|
||||
# 订阅主题
|
||||
topics = payload.get('topics', [])
|
||||
self.subscribe(client_id, *topics)
|
||||
|
||||
elif event_type == 'unsubscribe':
|
||||
# 取消订阅
|
||||
topics = payload.get('topics', [])
|
||||
self.unsubscribe(client_id, *topics)
|
||||
|
||||
elif event_type == 'ping':
|
||||
# 心跳响应
|
||||
self.send_to_client(client_id, 'pong', {'timestamp': time.time()})
|
||||
|
||||
elif event_type == 'status':
|
||||
# 请求状态 - 返回服务器和客户端状态
|
||||
response_data = {'status': 'ok'}
|
||||
|
||||
if payload.get('monitor'):
|
||||
# 获取实时监控数据
|
||||
try:
|
||||
if handler.monitor:
|
||||
response_data['monitor'] = handler.monitor.get_realtime_stats()
|
||||
else:
|
||||
# 回退到直接使用 psutil
|
||||
import psutil
|
||||
response_data['monitor'] = {
|
||||
'timestamp': datetime.now().isoformat(),
|
||||
'cpu': {
|
||||
'percent': psutil.cpu_percent(interval=0.1),
|
||||
'count': psutil.cpu_count()
|
||||
},
|
||||
'memory': psutil.virtual_memory()._asdict(),
|
||||
'disk': psutil.disk_usage(handler.config.get('base_dir', './downloads'))._asdict()
|
||||
}
|
||||
except Exception as e:
|
||||
response_data['monitor'] = {'error': str(e)}
|
||||
|
||||
if payload.get('ws'):
|
||||
# 获取 WebSocket 统计
|
||||
response_data['ws'] = self.get_stats()
|
||||
|
||||
if payload.get('client'):
|
||||
# 获取当前客户端信息
|
||||
client = self.clients.get(client_id)
|
||||
if client:
|
||||
response_data['client'] = {
|
||||
'client_id': client.client_id,
|
||||
'connected_at': client.connected_at,
|
||||
'last_activity': client.last_activity,
|
||||
'topics': list(client.topics),
|
||||
'metadata': client.metadata
|
||||
}
|
||||
|
||||
self.send_to_client(client_id, 'status:response', response_data)
|
||||
|
||||
elif event_type == 'sync':
|
||||
# 同步相关操作
|
||||
sync_action = payload.get('action', 'status')
|
||||
|
||||
if sync_action == 'status':
|
||||
# 获取同步状态
|
||||
response_data = {'action': 'status'}
|
||||
try:
|
||||
from core.sync_scheduler import SyncScheduler
|
||||
scheduler = SyncScheduler()
|
||||
response_data['sync_status'] = scheduler.get_status() if hasattr(scheduler, 'get_status') else {'message': 'sync scheduler running'}
|
||||
except Exception as e:
|
||||
response_data['sync_status'] = {'error': str(e)}
|
||||
|
||||
self.send_to_client(client_id, 'sync:response', response_data)
|
||||
|
||||
elif sync_action == 'list':
|
||||
# 获取同步任务列表
|
||||
response_data = {'action': 'list'}
|
||||
self.send_to_client(client_id, 'sync:response', response_data)
|
||||
|
||||
elif sync_action == 'trigger':
|
||||
# 触发手动同步
|
||||
response_data = {'action': 'trigger', 'status': 'pending'}
|
||||
self.send_to_client(client_id, 'sync:response', response_data)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"处理WebSocket消息失败: {e}")
|
||||
Reference in New Issue
Block a user