EG/plugins/user/realtime_communication/protocol/protocol_handler.py
2025-12-12 16:16:15 +08:00

658 lines
22 KiB
Python

"""
协议处理器模块
处理各种通信协议和消息格式
"""
import time
import json
import zlib
import base64
from typing import Dict, Any, List, Optional
class ProtocolHandler:
"""
协议处理器
处理各种通信协议和消息格式
"""
def __init__(self, plugin):
"""
初始化协议处理器
Args:
plugin: 实时通信插件实例
"""
self.plugin = plugin
self.enabled = False
self.initialized = False
# 协议配置
self.protocol_config = {
"supported_protocols": ["websocket", "json", "binary"],
"default_protocol": "json",
"enable_compression": True,
"compression_threshold": 1024, # 1KB
"enable_encryption": False,
"enable_message_validation": True,
"max_message_size": 65536, # 64KB
"enable_fragmentation": True,
"fragment_size": 8192, # 8KB
"enable_checksum": True
}
# 协议状态
self.protocol_state = {
"active_protocols": [],
"compressed_messages": 0,
"encrypted_messages": 0,
"fragmented_messages": 0
}
# 协议统计
self.protocol_stats = {
"messages_encoded": 0,
"messages_decoded": 0,
"bytes_encoded": 0,
"bytes_decoded": 0,
"compression_savings": 0,
"protocol_errors": 0
}
# 协议处理器映射
self.protocol_handlers = {
"json": {
"encode": self._encode_json_message,
"decode": self._decode_json_message
},
"binary": {
"encode": self._encode_binary_message,
"decode": self._decode_binary_message
}
}
# 回调函数
self.protocol_callbacks = {
"message_encoded": [],
"message_decoded": [],
"compression_applied": [],
"encryption_applied": [],
"protocol_error": []
}
# 时间戳记录
self.last_compression = 0.0
self.last_encryption = 0.0
print("✓ 协议处理器已创建")
def initialize(self) -> bool:
"""
初始化协议处理器
Returns:
是否初始化成功
"""
try:
print("正在初始化协议处理器...")
self.initialized = True
print("✓ 协议处理器初始化完成")
return True
except Exception as e:
print(f"✗ 协议处理器初始化失败: {e}")
self.protocol_stats["protocol_errors"] += 1
import traceback
traceback.print_exc()
return False
def enable(self) -> bool:
"""
启用协议处理器
Returns:
是否启用成功
"""
try:
if not self.initialized:
print("✗ 协议处理器未初始化")
return False
self.enabled = True
print("✓ 协议处理器已启用")
return True
except Exception as e:
print(f"✗ 协议处理器启用失败: {e}")
self.protocol_stats["protocol_errors"] += 1
import traceback
traceback.print_exc()
return False
def disable(self):
"""禁用协议处理器"""
try:
self.enabled = False
print("✓ 协议处理器已禁用")
except Exception as e:
print(f"✗ 协议处理器禁用失败: {e}")
self.protocol_stats["protocol_errors"] += 1
import traceback
traceback.print_exc()
def finalize(self):
"""清理协议处理器资源"""
try:
# 禁用协议处理器
if self.enabled:
self.disable()
# 清理回调
self.protocol_callbacks.clear()
self.initialized = False
print("✓ 协议处理器资源已清理")
except Exception as e:
print(f"✗ 协议处理器资源清理失败: {e}")
import traceback
traceback.print_exc()
def update(self, dt: float):
"""
更新协议处理器状态
Args:
dt: 时间增量(秒)
"""
try:
if not self.enabled:
return
# 协议处理器不需要频繁更新
pass
except Exception as e:
print(f"✗ 协议处理器更新失败: {e}")
self.protocol_stats["protocol_errors"] += 1
import traceback
traceback.print_exc()
def encode_message(self, message: Dict[str, Any], protocol: str = None) -> Optional[bytes]:
"""
编码消息
Args:
message: 消息数据
protocol: 协议类型
Returns:
编码后的消息或None
"""
try:
if not self.enabled:
return None
# 使用默认协议或指定协议
if protocol is None:
protocol = self.protocol_config["default_protocol"]
# 验证协议支持
if protocol not in self.protocol_handlers:
print(f"✗ 不支持的协议: {protocol}")
self.protocol_stats["protocol_errors"] += 1
return None
# 验证消息大小
if self.protocol_config["enable_message_validation"]:
message_size = len(json.dumps(message, ensure_ascii=False).encode('utf-8'))
if message_size > self.protocol_config["max_message_size"]:
print(f"✗ 消息过大: {message_size} bytes")
self.protocol_stats["protocol_errors"] += 1
return None
# 获取协议处理器
encoder = self.protocol_handlers[protocol]["encode"]
# 编码消息
encoded_message = encoder(message)
if encoded_message is None:
return None
# 更新统计
self.protocol_stats["messages_encoded"] += 1
self.protocol_stats["bytes_encoded"] += len(encoded_message)
# 触发消息编码回调
self._trigger_protocol_callback("message_encoded", {
"protocol": protocol,
"message_size": len(encoded_message),
"timestamp": time.time()
})
return encoded_message
except Exception as e:
print(f"✗ 消息编码失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return None
def decode_message(self, encoded_message: bytes, protocol: str = None) -> Optional[Dict[str, Any]]:
"""
解码消息
Args:
encoded_message: 编码后的消息
protocol: 协议类型
Returns:
解码后的消息或None
"""
try:
if not self.enabled:
return None
# 使用默认协议或指定协议
if protocol is None:
protocol = self.protocol_config["default_protocol"]
# 验证协议支持
if protocol not in self.protocol_handlers:
print(f"✗ 不支持的协议: {protocol}")
self.protocol_stats["protocol_errors"] += 1
return None
# 获取协议处理器
decoder = self.protocol_handlers[protocol]["decode"]
# 解码消息
decoded_message = decoder(encoded_message)
if decoded_message is None:
return None
# 更新统计
self.protocol_stats["messages_decoded"] += 1
self.protocol_stats["bytes_decoded"] += len(encoded_message)
# 触发消息解码回调
self._trigger_protocol_callback("message_decoded", {
"protocol": protocol,
"message_size": len(encoded_message),
"timestamp": time.time()
})
return decoded_message
except Exception as e:
print(f"✗ 消息解码失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return None
def _encode_json_message(self, message: Dict[str, Any]) -> Optional[bytes]:
"""
JSON消息编码
Args:
message: 消息数据
Returns:
编码后的消息或None
"""
try:
# 序列化为JSON
json_string = json.dumps(message, ensure_ascii=False)
encoded_message = json_string.encode('utf-8')
# 检查是否需要压缩
if (self.protocol_config["enable_compression"] and
len(encoded_message) > self.protocol_config["compression_threshold"]):
compressed_message = zlib.compress(encoded_message)
if len(compressed_message) < len(encoded_message):
# 使用压缩版本
encoded_message = compressed_message
self.protocol_state["compressed_messages"] += 1
self.protocol_stats["compression_savings"] += len(encoded_message) - len(compressed_message)
# 触发压缩应用回调
self._trigger_protocol_callback("compression_applied", {
"original_size": len(json_string.encode('utf-8')),
"compressed_size": len(compressed_message),
"timestamp": time.time()
})
return encoded_message
except Exception as e:
print(f"✗ JSON消息编码失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return None
def _decode_json_message(self, encoded_message: bytes) -> Optional[Dict[str, Any]]:
"""
JSON消息解码
Args:
encoded_message: 编码后的消息
Returns:
解码后的消息或None
"""
try:
# 尝试解压缩
try:
decompressed_message = zlib.decompress(encoded_message)
encoded_message = decompressed_message
except zlib.error:
# 不是压缩数据,使用原始数据
pass
# 解码为JSON
json_string = encoded_message.decode('utf-8')
decoded_message = json.loads(json_string)
return decoded_message
except Exception as e:
print(f"✗ JSON消息解码失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return None
def _encode_binary_message(self, message: Dict[str, Any]) -> Optional[bytes]:
"""
二进制消息编码
Args:
message: 消息数据
Returns:
编码后的消息或None
"""
try:
# 将消息转换为二进制格式(简化实现)
json_string = json.dumps(message, ensure_ascii=False)
binary_message = json_string.encode('utf-8')
# 添加简单的头部信息
header = len(binary_message).to_bytes(4, byteorder='big')
encoded_message = header + binary_message
# 检查是否需要压缩
if (self.protocol_config["enable_compression"] and
len(encoded_message) > self.protocol_config["compression_threshold"]):
compressed_message = zlib.compress(encoded_message)
if len(compressed_message) < len(encoded_message):
# 使用压缩版本
encoded_message = compressed_message
self.protocol_state["compressed_messages"] += 1
self.protocol_stats["compression_savings"] += len(encoded_message) - len(compressed_message)
# 触发压缩应用回调
self._trigger_protocol_callback("compression_applied", {
"original_size": len(header + binary_message),
"compressed_size": len(compressed_message),
"timestamp": time.time()
})
return encoded_message
except Exception as e:
print(f"✗ 二进制消息编码失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return None
def _decode_binary_message(self, encoded_message: bytes) -> Optional[Dict[str, Any]]:
"""
二进制消息解码
Args:
encoded_message: 编码后的消息
Returns:
解码后的消息或None
"""
try:
# 尝试解压缩
try:
decompressed_message = zlib.decompress(encoded_message)
encoded_message = decompressed_message
except zlib.error:
# 不是压缩数据,使用原始数据
pass
# 解析头部信息
if len(encoded_message) < 4:
raise ValueError("消息太短")
message_length = int.from_bytes(encoded_message[:4], byteorder='big')
binary_message = encoded_message[4:4 + message_length]
# 解码为JSON
json_string = binary_message.decode('utf-8')
decoded_message = json.loads(json_string)
return decoded_message
except Exception as e:
print(f"✗ 二进制消息解码失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return None
def fragment_message(self, message: bytes, fragment_size: int = None) -> List[bytes]:
"""
分片消息
Args:
message: 消息数据
fragment_size: 分片大小
Returns:
分片列表
"""
try:
if not self.protocol_config["enable_fragmentation"]:
return [message]
if fragment_size is None:
fragment_size = self.protocol_config["fragment_size"]
fragments = []
for i in range(0, len(message), fragment_size):
fragment = message[i:i + fragment_size]
fragments.append(fragment)
self.protocol_state["fragmented_messages"] += 1
return fragments
except Exception as e:
print(f"✗ 消息分片失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return [message]
def reassemble_message(self, fragments: List[bytes]) -> Optional[bytes]:
"""
重组消息
Args:
fragments: 分片列表
Returns:
重组后的消息或None
"""
try:
if not fragments:
return None
reassembled_message = b''.join(fragments)
return reassembled_message
except Exception as e:
print(f"✗ 消息重组失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return None
def add_protocol_handler(self, protocol: str, encode_func: callable, decode_func: callable) -> bool:
"""
添加协议处理器
Args:
protocol: 协议名称
encode_func: 编码函数
decode_func: 解码函数
Returns:
是否添加成功
"""
try:
self.protocol_handlers[protocol] = {
"encode": encode_func,
"decode": decode_func
}
if protocol not in self.protocol_config["supported_protocols"]:
self.protocol_config["supported_protocols"].append(protocol)
print(f"✓ 协议处理器已添加: {protocol}")
return True
except Exception as e:
print(f"✗ 协议处理器添加失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return False
def remove_protocol_handler(self, protocol: str) -> bool:
"""
移除协议处理器
Args:
protocol: 协议名称
Returns:
是否移除成功
"""
try:
if protocol in self.protocol_handlers:
del self.protocol_handlers[protocol]
if protocol in self.protocol_config["supported_protocols"]:
self.protocol_config["supported_protocols"].remove(protocol)
print(f"✓ 协议处理器已移除: {protocol}")
return True
else:
print(f"✗ 协议处理器不存在: {protocol}")
return False
except Exception as e:
print(f"✗ 协议处理器移除失败: {e}")
self.protocol_stats["protocol_errors"] += 1
return False
def get_protocol_stats(self) -> Dict[str, Any]:
"""
获取协议统计信息
Returns:
协议统计字典
"""
return {
"state": self.protocol_state.copy(),
"stats": self.protocol_stats.copy(),
"config": self.protocol_config.copy()
}
def reset_stats(self):
"""重置协议统计信息"""
try:
self.protocol_stats = {
"messages_encoded": 0,
"messages_decoded": 0,
"bytes_encoded": 0,
"bytes_decoded": 0,
"compression_savings": 0,
"protocol_errors": 0
}
print("✓ 协议统计信息已重置")
except Exception as e:
print(f"✗ 协议统计信息重置失败: {e}")
def set_protocol_config(self, config: Dict[str, Any]) -> bool:
"""
设置协议配置
Args:
config: 协议配置字典
Returns:
是否设置成功
"""
try:
self.protocol_config.update(config)
print(f"✓ 协议配置已更新: {self.protocol_config}")
return True
except Exception as e:
print(f"✗ 协议配置设置失败: {e}")
return False
def get_protocol_config(self) -> Dict[str, Any]:
"""
获取协议配置
Returns:
协议配置字典
"""
return self.protocol_config.copy()
def _trigger_protocol_callback(self, callback_type: str, data: Dict[str, Any]):
"""
触发协议回调
Args:
callback_type: 回调类型
data: 回调数据
"""
try:
if callback_type in self.protocol_callbacks:
for callback in self.protocol_callbacks[callback_type]:
try:
callback(data)
except Exception as e:
print(f"✗ 协议回调执行失败: {callback_type} - {e}")
except Exception as e:
print(f"✗ 协议回调触发失败: {e}")
def register_protocol_callback(self, callback_type: str, callback: callable):
"""
注册协议回调
Args:
callback_type: 回调类型
callback: 回调函数
"""
try:
if callback_type in self.protocol_callbacks:
self.protocol_callbacks[callback_type].append(callback)
print(f"✓ 协议回调已注册: {callback_type}")
else:
print(f"✗ 无效的回调类型: {callback_type}")
except Exception as e:
print(f"✗ 协议回调注册失败: {e}")
def unregister_protocol_callback(self, callback_type: str, callback: callable):
"""
注销协议回调
Args:
callback_type: 回调类型
callback: 回调函数
"""
try:
if callback_type in self.protocol_callbacks:
if callback in self.protocol_callbacks[callback_type]:
self.protocol_callbacks[callback_type].remove(callback)
print(f"✓ 协议回调已注销: {callback_type}")
else:
print(f"✗ 无效的回调类型: {callback_type}")
except Exception as e:
print(f"✗ 协议回调注销失败: {e}")