Files
banban/talkingq-url/handlers/websocket_auth_handler.py
2026-03-24 15:04:36 +08:00

67 lines
2.7 KiB
Python

import json
import asyncio
import time
from fastapi import WebSocket
from services.device_auth_manager import device_auth_manager
from services.connection_manager import connection_manager
from utils.logger import session_logger
async def authenticate_websocket(websocket: WebSocket):
"""
处理设备认证流程
Args:
websocket: WebSocket连接
Returns:
tuple: (认证状态, 设备ID) - (是否认证成功, 设备ID)
"""
device_id = None
authenticated = False
auth_timeout = 10 # 10秒认证超时
auth_start_time = time.time()
while not authenticated and time.time() - auth_start_time < auth_timeout:
try:
message = await asyncio.wait_for(
websocket.receive(),
timeout=auth_timeout - (time.time() - auth_start_time)
)
if message["type"] == "websocket.receive" and "text" in message:
try:
auth_data = json.loads(message["text"])
if "device_id" in auth_data and "serial_number" in auth_data:
device_id = auth_data["device_id"]
serial_number = auth_data["serial_number"]
authenticated = await device_auth_manager.authenticate_device(
device_id, serial_number
)
if authenticated:
await connection_manager.add_connection(device_id, websocket)
session_logger.info(
device_id, "auth", f"设备 {device_id} 认证成功"
)
await websocket.send_text(
json.dumps({"status": "authenticated"})
)
else:
session_logger.warning(
device_id, "auth", f"设备 {device_id} 认证失败"
)
await websocket.send_text(
json.dumps(
{"status": "error", "message": "Authentication failed"}
)
)
except json.JSONDecodeError:
session_logger.warning("unknown", "auth", "收到的不是有效的JSON认证消息")
elif message["type"] == "websocket.receive" and "bytes" in message:
pass
except asyncio.TimeoutError:
session_logger.warning("unknown", "auth", "WebSocket认证超时")
break
return authenticated, device_id