274 lines
9.7 KiB
Python
274 lines
9.7 KiB
Python
|
import asyncio
|
|||
|
|
import websockets
|
|||
|
|
import json
|
|||
|
|
import ssl
|
|||
|
|
import secrets
|
|||
|
|
import aiohttp
|
|||
|
|
from aiohttp import web
|
|||
|
|
import yaml
|
|||
|
|
import logging
|
|||
|
|
from datetime import datetime
|
|||
|
|
from collections import defaultdict
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
# 配置日志
|
|||
|
|
logging.basicConfig(
|
|||
|
|
level=logging.INFO,
|
|||
|
|
format='%(asctime)s - %(levelname)s - %(message)s'
|
|||
|
|
)
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
class NotificationServer:
|
|||
|
|
def __init__(self, config_path='config.yml'):
|
|||
|
|
try:
|
|||
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|||
|
|
self.config = yaml.safe_load(f)
|
|||
|
|
except FileNotFoundError:
|
|||
|
|
# 如果配置文件不存在，创建默认配置
|
|||
|
|
self.config = {
|
|||
|
|
'websocket': {
|
|||
|
|
'host': '0.0.0.0',
|
|||
|
|
'port': 8080
|
|||
|
|
},
|
|||
|
|
'security': {
|
|||
|
|
'use_ssl': False,
|
|||
|
|
'ssl_cert': 'path/to/cert.pem',
|
|||
|
|
'ssl_key': 'path/to/key.pem'
|
|||
|
|
},
|
|||
|
|
'data': {
|
|||
|
|
'keys_file': '/app/data/registered_keys.json'
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
# 保存默认配置
|
|||
|
|
with open(config_path, 'w', encoding='utf-8') as f:
|
|||
|
|
yaml.safe_dump(self.config, f, allow_unicode=True)
|
|||
|
|
|
|||
|
|
# 存储客户端连接
|
|||
|
|
self.client_connections = defaultdict(set)
|
|||
|
|
# 从配置文件获取数据存储路径
|
|||
|
|
self.keys_file = self.config['data']['keys_file']
|
|||
|
|
self.registered_keys = {}
|
|||
|
|
self.load_registered_keys()
|
|||
|
|
|
|||
|
|
def load_registered_keys(self):
|
|||
|
|
"""从文件加载已注册的密钥"""
|
|||
|
|
try:
|
|||
|
|
if os.path.exists(self.keys_file):
|
|||
|
|
with open(self.keys_file, 'r', encoding='utf-8') as f:
|
|||
|
|
data = json.load(f)
|
|||
|
|
self.registered_keys = data
|
|||
|
|
logger.info(f"已加载 {len(self.registered_keys)} 个注册密钥")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"加载密钥文件失败: {e}")
|
|||
|
|
|
|||
|
|
def save_registered_keys(self):
|
|||
|
|
"""保存密钥到文件"""
|
|||
|
|
try:
|
|||
|
|
with open(self.keys_file, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(self.registered_keys, f, indent=2)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"保存密钥文件失败: {e}")
|
|||
|
|
|
|||
|
|
def generate_key(self):
|
|||
|
|
"""生成唯一的客户端密钥"""
|
|||
|
|
return secrets.token_urlsafe(32)
|
|||
|
|
|
|||
|
|
async def register_client(self, websocket):
|
|||
|
|
"""注册新客户端并生成密钥"""
|
|||
|
|
try:
|
|||
|
|
# 等待客户端发送注册请求
|
|||
|
|
msg = await websocket.recv()
|
|||
|
|
data = json.loads(msg)
|
|||
|
|
|
|||
|
|
if data.get('action') == 'register':
|
|||
|
|
device_name = data.get('device_name', 'unknown')
|
|||
|
|
key = self.generate_key()
|
|||
|
|
|
|||
|
|
client_info = {
|
|||
|
|
'device_name': device_name,
|
|||
|
|
'created_at': datetime.now().isoformat(),
|
|||
|
|
'last_seen': datetime.now().isoformat()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
self.registered_keys[key] = client_info
|
|||
|
|
self.client_connections[key].add(websocket)
|
|||
|
|
|
|||
|
|
# 发送密钥给客户端
|
|||
|
|
await websocket.send(json.dumps({
|
|||
|
|
'status': 'success',
|
|||
|
|
'key': key,
|
|||
|
|
'message': '注册成功'
|
|||
|
|
}))
|
|||
|
|
|
|||
|
|
logger.info(f"新客户端注册: {device_name}, key: {key[:8]}...")
|
|||
|
|
self.save_registered_keys() # 保存新注册的密钥
|
|||
|
|
return key
|
|||
|
|
|
|||
|
|
elif data.get('action') == 'connect':
|
|||
|
|
key = data.get('key')
|
|||
|
|
if key in self.registered_keys:
|
|||
|
|
self.client_connections[key].add(websocket)
|
|||
|
|
self.registered_keys[key]['last_seen'] = datetime.now().isoformat()
|
|||
|
|
await websocket.send(json.dumps({
|
|||
|
|
'status': 'success',
|
|||
|
|
'message': '连接成功'
|
|||
|
|
}))
|
|||
|
|
logger.info(f"客户端重连成功: {key[:8]}...")
|
|||
|
|
return key
|
|||
|
|
else:
|
|||
|
|
await websocket.send(json.dumps({
|
|||
|
|
'status': 'error',
|
|||
|
|
'message': '无效的密钥'
|
|||
|
|
}))
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"注册失败: {e}")
|
|||
|
|
await websocket.close()
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
async def handle_websocket(self, websocket):
|
|||
|
|
key = await self.register_client(websocket)
|
|||
|
|
if not key:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 添加心跳检测
|
|||
|
|
while True:
|
|||
|
|
try:
|
|||
|
|
message = await asyncio.wait_for(websocket.recv(), timeout=30)
|
|||
|
|
if message == "ping":
|
|||
|
|
await websocket.send("pong")
|
|||
|
|
except asyncio.TimeoutError:
|
|||
|
|
# 30秒没有收到消息，发送 ping
|
|||
|
|
try:
|
|||
|
|
await websocket.send("ping")
|
|||
|
|
await asyncio.wait_for(websocket.recv(), timeout=10)
|
|||
|
|
except:
|
|||
|
|
# ping 失败，关闭连接
|
|||
|
|
break
|
|||
|
|
except websockets.exceptions.ConnectionClosed:
|
|||
|
|
break
|
|||
|
|
finally:
|
|||
|
|
self.client_connections[key].discard(websocket)
|
|||
|
|
if not self.client_connections[key]:
|
|||
|
|
logger.info(f"客户端完全断开: {key[:8]}...")
|
|||
|
|
|
|||
|
|
async def send_notification(self, key, notification):
|
|||
|
|
"""向指定key的所有连接发送通知"""
|
|||
|
|
if key not in self.client_connections:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
connections = self.client_connections[key].copy()
|
|||
|
|
if not connections:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
for websocket in connections:
|
|||
|
|
try:
|
|||
|
|
await websocket.send(json.dumps(notification))
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"发送通知失败: {e}")
|
|||
|
|
self.client_connections[key].discard(websocket)
|
|||
|
|
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
async def handle_http_push(self, request):
|
|||
|
|
try:
|
|||
|
|
# 从URL路径中获取key
|
|||
|
|
key = request.match_info['key']
|
|||
|
|
|
|||
|
|
if key not in self.registered_keys:
|
|||
|
|
return web.Response(status=404, text='Invalid key')
|
|||
|
|
|
|||
|
|
data = await request.json()
|
|||
|
|
if not data.get('title') or not data.get('message'):
|
|||
|
|
return web.Response(status=400, text='Missing title or message')
|
|||
|
|
|
|||
|
|
notification = {
|
|||
|
|
'title': data['title'],
|
|||
|
|
'message': data['message'],
|
|||
|
|
'type': data.get('type', 'info'),
|
|||
|
|
'timestamp': datetime.now().isoformat()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
success = await self.send_notification(key, notification)
|
|||
|
|
if success:
|
|||
|
|
return web.Response(text='Notification sent')
|
|||
|
|
else:
|
|||
|
|
return web.Response(status=404, text='No active connections')
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"处理HTTP请求失败: {e}")
|
|||
|
|
return web.Response(status=500, text='Internal server error')
|
|||
|
|
|
|||
|
|
async def start(self):
|
|||
|
|
# 设置SSL
|
|||
|
|
ssl_context = None
|
|||
|
|
if self.config['security']['use_ssl']:
|
|||
|
|
ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
|
|||
|
|
ssl_context.load_cert_chain(
|
|||
|
|
self.config['security']['ssl_cert'],
|
|||
|
|
self.config['security']['ssl_key']
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 创建 aiohttp 应用
|
|||
|
|
app = web.Application()
|
|||
|
|
|
|||
|
|
# 添加 WebSocket 处理
|
|||
|
|
async def websocket_handler(request):
|
|||
|
|
ws = web.WebSocketResponse()
|
|||
|
|
await ws.prepare(request)
|
|||
|
|
|
|||
|
|
# 包装 aiohttp WebSocketResponse 为 websockets 兼容的接口
|
|||
|
|
class WebSocketWrapper:
|
|||
|
|
async def send(self, message):
|
|||
|
|
await ws.send_str(message)
|
|||
|
|
|
|||
|
|
async def recv(self):
|
|||
|
|
msg = await ws.receive()
|
|||
|
|
if msg.type == web.WSMsgType.TEXT:
|
|||
|
|
return msg.data
|
|||
|
|
elif msg.type == web.WSMsgType.CLOSE:
|
|||
|
|
raise websockets.exceptions.ConnectionClosed(None, None)
|
|||
|
|
|
|||
|
|
async def close(self):
|
|||
|
|
await ws.close()
|
|||
|
|
|
|||
|
|
wrapped_ws = WebSocketWrapper()
|
|||
|
|
await self.handle_websocket(wrapped_ws)
|
|||
|
|
return ws
|
|||
|
|
|
|||
|
|
# 注册路由
|
|||
|
|
app.router.add_get('/ws', websocket_handler) # WebSocket 端点
|
|||
|
|
app.router.add_post('/push/{key}', self.handle_http_push) # HTTP 推送端点
|
|||
|
|
|
|||
|
|
# 启动服务器
|
|||
|
|
runner = web.AppRunner(app)
|
|||
|
|
await runner.setup()
|
|||
|
|
|
|||
|
|
# 根据是否使用 SSL 创建不同的站点
|
|||
|
|
if ssl_context:
|
|||
|
|
site = web.SSLSite(
|
|||
|
|
runner,
|
|||
|
|
self.config['websocket']['host'],
|
|||
|
|
self.config['websocket']['port'],
|
|||
|
|
ssl_context=ssl_context
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
site = web.TCPSite(
|
|||
|
|
runner,
|
|||
|
|
self.config['websocket']['host'],
|
|||
|
|
self.config['websocket']['port']
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
await site.start()
|
|||
|
|
|
|||
|
|
logger.info(f"服务器运行在 {self.config['websocket']['host']}:{self.config['websocket']['port']}")
|
|||
|
|
logger.info("WebSocket 路径: /ws")
|
|||
|
|
logger.info("HTTP 推送路径: /push/{key}")
|
|||
|
|
|
|||
|
|
await asyncio.Future() # 持续运行
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
server = NotificationServer()
|
|||
|
|
asyncio.run(server.start())
|