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()) |