优化代码风格

ws/ws.py

@@ -2,38 +2,81 @@ import asyncio
 import datetime
 import json
 import os
+import base64
 from contextlib import asynccontextmanager
 from typing import Dict, Optional
+
+import cv2
+import numpy as np
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import pad, unpad
+from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI
+
 from service.device_service import update_online_status_by_ip, increment_alarm_count_by_ip
 from service.device_action_service import add_device_action
 from schema.device_action_schema import DeviceActionCreate
 from core.all import detect, load_model
 
-import cv2
+# -------------------------- 1. AES 加密解密工具（固定密钥）--------------------------
-import numpy as np
+AES_SECRET_KEY = b"jr1vA6tfWMHOYi6UXw67UuO6fdak2rMa"  # 约定密钥（32字节）
-from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI
+AES_BLOCK_SIZE = 16  # AES固定块大小
 
-# 配置常量
+
+def aes_encrypt(plaintext: str) -> dict:
+    """AES-CBC加密：返回{密文, IV, 算法标识}（均Base64编码）"""
+    try:
+        iv = os.urandom(AES_BLOCK_SIZE)  # 随机IV（16字节）
+        cipher = AES.new(AES_SECRET_KEY, AES.MODE_CBC, iv)
+        # 明文填充+加密+Base64编码
+        padded_plaintext = pad(plaintext.encode("utf-8"), AES_BLOCK_SIZE)
+        ciphertext = base64.b64encode(cipher.encrypt(padded_plaintext)).decode("utf-8")
+        iv_base64 = base64.b64encode(iv).decode("utf-8")
+        return {
+            "ciphertext": ciphertext,
+            "iv": iv_base64,
+            "algorithm": "AES-CBC"
+        }
+    except Exception as e:
+        raise Exception(f"AES加密失败: {str(e)}") from e
+
+
+def aes_decrypt(encrypted_dict: dict) -> str:
+    """AES-CBC解密：输入加密字典，返回原始文本"""
+    try:
+        # Base64解码密文和IV
+        ciphertext = base64.b64decode(encrypted_dict["ciphertext"])
+        iv = base64.b64decode(encrypted_dict["iv"])
+        # 解密+去除填充
+        cipher = AES.new(AES_SECRET_KEY, AES.MODE_CBC, iv)
+        decrypted = unpad(cipher.decrypt(ciphertext), AES_BLOCK_SIZE).decode("utf-8")
+        return decrypted
+    except Exception as e:
+        raise Exception(f"AES解密失败: {str(e)}") from e
+
+
+# -------------------------- 2. 配置常量（保持原有）--------------------------
 HEARTBEAT_INTERVAL = 30  # 心跳检查间隔（秒）
 HEARTBEAT_TIMEOUT = 600  # 客户端超时阈值（秒）
 WS_ENDPOINT = "/ws"  # WebSocket端点路径
 FRAME_QUEUE_SIZE = 1  # 帧队列大小限制
 
 
-# 工具函数: 获取格式化时间字符串
+# -------------------------- 3. 工具函数（保持原有）--------------------------
 def get_current_time_str() -> str:
+    """获取格式化时间字符串"""
     return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
 
 def get_current_time_file_str() -> str:
+    """获取文件命名用时间字符串"""
     return datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
 
 
-# 客户端连接封装
+# -------------------------- 4. 客户端连接封装（新增消息加密）--------------------------
 class ClientConnection:
     def __init__(self, websocket: WebSocket, client_ip: str):
         self.websocket = websocket
-        self.client_ip = client_ip  # 已初始化客户端IP，用于传递给detect
+        self.client_ip = client_ip
         self.last_heartbeat = datetime.datetime.now()
         self.frame_queue = asyncio.Queue(maxsize=FRAME_QUEUE_SIZE)
         self.consumer_task: Optional[asyncio.Task] = None
@@ -53,96 +96,89 @@ class ClientConnection:
         return self.consumer_task
 
     async def send_frame_permit(self):
-        """发送帧发送许可信号"""
+        """发送加密的帧许可信号"""
         try:
+            # 1. 构建原始消息
             frame_permit_msg = {
                 "type": "frame",
                 "timestamp": get_current_time_str(),
                 "client_ip": self.client_ip
             }
-            await self.websocket.send_json(frame_permit_msg)
+            # 2. AES加密消息
-            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 已发送帧发送许可信号")
+            encrypted_msg = aes_encrypt(json.dumps(frame_permit_msg))
+            # 3. 发送加密消息
+            await self.websocket.send_json(encrypted_msg)
+            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 已发送加密帧许可")
         except Exception as e:
-            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧许可信号发送失败 - {str(e)}")
+            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧许可加密/发送失败 - {str(e)}")
 
     async def consume_frames(self) -> None:
         """消费队列中的帧并处理"""
         try:
             while True:
-                # 取出帧并立即发送下一帧许可
                 frame_data = await self.frame_queue.get()
-                await self.send_frame_permit()
+                await self.send_frame_permit()  # 发送下一帧许可
-
                 try:
                     await self.process_frame(frame_data)
                 finally:
                     self.frame_queue.task_done()
-
         except asyncio.CancelledError:
             print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧消费任务已取消")
         except Exception as e:
-            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧消费逻辑错误 - {str(e)}")
+            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧消费错误 - {str(e)}")
 
     async def process_frame(self, frame_data: bytes) -> None:
-        """处理单帧图像数据（核心修改：detect函数传入 client_ip + img 双参数）"""
+        """处理单帧图像（含加密危险通知）"""
         # 二进制转OpenCV图像
         nparr = np.frombuffer(frame_data, np.uint8)
         img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
         if img is None:
-            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 无法解析图像数据")
+            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 无法解析图像")
             return
 
         try:
-            # -------------------------- 核心修改：按要求传入参数（1.client_ip 2.img） --------------------------
+            # 调用检测函数（client_ip + img 双参数）
-            # detect函数参数顺序：第一个为client_ip，第二个为图像数据img
-            # 保持返回值解包（是否违规, 结果数据, 检测器类型）不变
             has_violation, data, detector_type = await asyncio.to_thread(
-                detect,                  # 调用检测函数
+                detect, self.client_ip, img
-                self.client_ip,          # 第一个参数：客户端IP（新增，按需求顺序）
-                img                      # 第二个参数：图像数据（原参数，调整顺序）
             )
-            # -------------------------------------------------------------------------------------
+            print(
+                f"[{get_current_time_str()}] 客户端{self.client_ip}: 检测结果 - 违规: {has_violation}, 类型: {detector_type}")
 
-            # 打印检测结果（包含客户端IP，与传入参数对应）
+            # 处理违规逻辑（发送加密危险通知）
-            print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 检测结果 - "
-                  f"违规: {has_violation}, 类型: {detector_type}, 数据: {data}")
-
-            # 处理违规逻辑（逻辑不变，基于detect返回结果执行）
             if has_violation:
-                print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 检测到违规 - "
+                print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 检测到违规 - {data}")
-                      f"类型: {detector_type}, 详情: {data}")
-
                 # 违规次数+1
                 try:
                     await asyncio.to_thread(increment_alarm_count_by_ip, self.client_ip)
-                    print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 违规次数已+1")
+                    print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 违规次数+1")
                 except Exception as e:
                     print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 违规次数更新失败 - {str(e)}")
 
-                # 发送危险通知
+                # 1. 构建原始危险通知
                 danger_msg = {
                     "type": "danger",
                     "timestamp": get_current_time_str(),
                     "client_ip": self.client_ip,
                     "detail": data
                 }
-                await self.websocket.send_json(danger_msg)
+                # 2. AES加密通知
-            else:
+                encrypted_danger_msg = aes_encrypt(json.dumps(danger_msg))
-                print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 未检测到违规")
+                # 3. 发送加密通知
-
+                await self.websocket.send_json(encrypted_danger_msg)
         except Exception as e:
             print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 图像处理错误 - {str(e)}")
 
 
-# 全局状态管理
+# -------------------------- 5. 全局状态与心跳管理（保持原有）--------------------------
 connected_clients: Dict[str, ClientConnection] = {}
 heartbeat_task: Optional[asyncio.Task] = None
 
 
-# 心跳检查任务
 async def heartbeat_checker():
+    """全局心跳检查任务"""
     while True:
         current_time = get_current_time_str()
+        # 筛选超时客户端
         timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()]
 
         if timeout_ips:
@@ -150,18 +186,17 @@ async def heartbeat_checker():
             for ip in timeout_ips:
                 try:
                     conn = connected_clients[ip]
+                    # 取消消费任务+关闭连接
                     if conn.consumer_task and not conn.consumer_task.done():
                         conn.consumer_task.cancel()
                     await conn.websocket.close(code=1008, reason="心跳超时")
-
                     # 标记离线
-                    try:
                     await asyncio.to_thread(update_online_status_by_ip, ip, 0)
                     action_data = DeviceActionCreate(client_ip=ip, action=0)
                     await asyncio.to_thread(add_device_action, action_data)
-                        print(f"[{current_time}] 客户端{ip}: 已标记为离线并记录操作")
+                    print(f"[{current_time}] 客户端{ip}: 已标记为离线")
                 except Exception as e:
-                        print(f"[{current_time}] 客户端{ip}: 离线状态更新失败 - {str(e)}")
+                    print(f"[{current_time}] 客户端{ip}: 离线处理失败 - {str(e)}")
                 finally:
                     connected_clients.pop(ip, None)
         else:
@@ -170,87 +205,108 @@ async def heartbeat_checker():
         await asyncio.sleep(HEARTBEAT_INTERVAL)
 
 
-# 应用生命周期管理
+# -------------------------- 6. 消息处理工具（新增消息解密）--------------------------
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    global heartbeat_task
-    heartbeat_task = asyncio.create_task(heartbeat_checker())
-    print(f"[{get_current_time_str()}] 全局心跳检查任务启动（任务ID: {id(heartbeat_task)}）")
-    yield
-    if heartbeat_task and not heartbeat_task.done():
-        heartbeat_task.cancel()
-        try:
-            await heartbeat_task
-            print(f"[{get_current_time_str()}] 全局心跳检查任务已取消")
-        except asyncio.CancelledError:
-            pass
-
-
-# 消息处理工具函数
 async def send_heartbeat_ack(conn: ClientConnection):
+    """发送加密的心跳确认"""
     try:
+        # 1. 构建原始心跳确认
         heartbeat_ack_msg = {
             "type": "heart",
             "timestamp": get_current_time_str(),
             "client_ip": conn.client_ip
         }
-        await conn.websocket.send_json(heartbeat_ack_msg)
+        # 2. AES加密
-        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 已发送心跳确认")
+        encrypted_msg = aes_encrypt(json.dumps(heartbeat_ack_msg))
+        # 3. 发送
+        await conn.websocket.send_json(encrypted_msg)
+        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 已发送加密心跳确认")
         return True
     except Exception as e:
         connected_clients.pop(conn.client_ip, None)
-        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 心跳确认发送失败 - {str(e)}")
+        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 心跳确认失败 - {str(e)}")
         return False
 
 
 async def handle_text_msg(conn: ClientConnection, text: str):
+    """处理加密的文本消息（如心跳）"""
     try:
-        msg = json.loads(text)
+        # 1. 解析加密字典
+        encrypted_dict = json.loads(text)
+        # 2. AES解密
+        decrypted_text = aes_decrypt(encrypted_dict)
+        # 3. 解析业务消息
+        msg = json.loads(decrypted_text)
+
         if msg.get("type") == "heart":
             conn.update_heartbeat()
             await send_heartbeat_ack(conn)
         else:
-            print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 未知文本消息类型（{msg.get('type')}）")
+            print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 未知文本类型（{msg.get('type')}）")
     except json.JSONDecodeError:
-        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 无效JSON文本消息")
+        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 无效JSON格式")
+    except Exception as e:
+        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 文本消息解密失败 - {str(e)}")
 
 
-async def handle_binary_msg(conn: ClientConnection, data: bytes):
+async def handle_binary_msg(conn: ClientConnection, data: str):
+    """处理加密的图像消息（客户端需先转Base64+加密）"""
     try:
-        conn.frame_queue.put_nowait(data)
+        # 1. 解密得到Base64编码的图像
-        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 图像数据（{len(data)}字节）已加入队列")
+        encrypted_dict = json.loads(data)
+        decrypted_base64 = aes_decrypt(encrypted_dict)
+        # 2. Base64解码为二进制图像
+        frame_data = base64.b64decode(decrypted_base64)
+        # 3. 加入帧队列
+        conn.frame_queue.put_nowait(frame_data)
+        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 解密后图像（{len(frame_data)}字节）入队")
     except asyncio.QueueFull:
-        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 帧队列已满、丢弃当前图像数据")
+        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 帧队列已满，丢弃数据")
+    except Exception as e:
+        print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 图像消息解密失败 - {str(e)}")
 
 
-# WebSocket路由配置
+# -------------------------- 7. WebSocket路由与生命周期（保持原有结构）--------------------------
 ws_router = APIRouter()
 
 
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """应用生命周期：启动/停止心跳任务"""
+    global heartbeat_task
+    heartbeat_task = asyncio.create_task(heartbeat_checker())
+    print(f"[{get_current_time_str()}] 心跳检查任务启动（ID: {id(heartbeat_task)}）")
+    yield
+    # 关闭时清理
+    if heartbeat_task and not heartbeat_task.done():
+        heartbeat_task.cancel()
+        await heartbeat_task
+        print(f"[{get_current_time_str()}] 心跳检查任务已取消")
+
+
 @ws_router.websocket(WS_ENDPOINT)
 async def websocket_endpoint(websocket: WebSocket):
-    load_model()  # 加载检测模型（仅在连接建立时加载一次，避免重复加载）
+    """WebSocket连接处理入口"""
+    load_model()  # 加载检测模型（仅一次）
     await websocket.accept()
     client_ip = websocket.client.host if websocket.client else "unknown_ip"
     current_time = get_current_time_str()
-    print(f"[{current_time}] 客户端{client_ip}: WebSocket连接已建立")
+    print(f"[{current_time}] 客户端{client_ip}: 连接已建立")
-
     is_online_updated = False
 
     try:
-        # 处理重复连接（同一IP断开旧连接）
+        # 处理重复连接（关闭旧连接）
         if client_ip in connected_clients:
             old_conn = connected_clients[client_ip]
             if old_conn.consumer_task and not old_conn.consumer_task.done():
                 old_conn.consumer_task.cancel()
-            await old_conn.websocket.close(code=1008, reason="同一IP新连接建立")
+            await old_conn.websocket.close(code=1008, reason="新连接建立")
             connected_clients.pop(client_ip)
             print(f"[{current_time}] 客户端{client_ip}: 已关闭旧连接")
 
-        # 注册新连接（绑定client_ip和WebSocket）
+        # 注册新连接
         new_conn = ClientConnection(websocket, client_ip)
         connected_clients[client_ip] = new_conn
-        new_conn.start_consumer()  # 启动帧消费任务
+        new_conn.start_consumer()  # 启动帧消费
         await new_conn.send_frame_permit()  # 发送首次帧许可
 
         # 标记客户端上线
@@ -258,41 +314,42 @@ async def websocket_endpoint(websocket: WebSocket):
             await asyncio.to_thread(update_online_status_by_ip, client_ip, 1)
             action_data = DeviceActionCreate(client_ip=client_ip, action=1)
             await asyncio.to_thread(add_device_action, action_data)
-            print(f"[{current_time}] 客户端{client_ip}: 已标记为在线并记录操作")
+            print(f"[{current_time}] 客户端{client_ip}: 已标记为在线")
             is_online_updated = True
         except Exception as e:
             print(f"[{current_time}] 客户端{client_ip}: 上线状态更新失败 - {str(e)}")
 
-        print(f"[{current_time}] 客户端{client_ip}: 新连接注册成功、在线数: {len(connected_clients)}")
+        print(f"[{current_time}] 客户端{client_ip}: 连接注册成功，在线数: {len(connected_clients)}")
 
-        # 消息循环（持续接收客户端消息）
+        # 消息循环（接收客户端消息）
         while True:
             data = await websocket.receive()
             if "text" in data:
+                # 处理加密文本消息（心跳、客户端指令）
                 await handle_text_msg(new_conn, data["text"])
             elif "bytes" in data:
-                await handle_binary_msg(new_conn, data["bytes"])
+                # 兼容客户端发送二进制：先转Base64再处理
+                base64_data = base64.b64encode(data["bytes"]).decode("utf-8")
+                await handle_binary_msg(new_conn, base64_data)
 
     except WebSocketDisconnect as e:
         print(f"[{get_current_time_str()}] 客户端{client_ip}: 主动断开连接（代码: {e.code}）")
     except Exception as e:
         print(f"[{get_current_time_str()}] 客户端{client_ip}: 连接异常 - {str(e)[:50]}")
     finally:
-        # 清理资源（断开后标记离线+删除连接）
+        # 清理资源（断开后处理）
         if client_ip in connected_clients:
             conn = connected_clients[client_ip]
             if conn.consumer_task and not conn.consumer_task.done():
                 conn.consumer_task.cancel()
-
+            # 仅上线成功的客户端，才标记离线
-            # 仅当上线状态更新成功时，才执行离线更新
             if is_online_updated:
                 try:
                     await asyncio.to_thread(update_online_status_by_ip, client_ip, 0)
                     action_data = DeviceActionCreate(client_ip=client_ip, action=0)
                     await asyncio.to_thread(add_device_action, action_data)
-                    print(f"[{get_current_time_str()}] 客户端{client_ip}: 断开后已标记为离线")
+                    print(f"[{get_current_time_str()}] 客户端{client_ip}: 断开后标记为离线")
                 except Exception as e:
-                    print(f"[{get_current_time_str()}] 客户端{client_ip}: 断开后离线更新失败 - {str(e)}")
+                    print(f"[{get_current_time_str()}] 客户端{client_ip}: 离线更新失败 - {str(e)}")
-
             connected_clients.pop(client_ip, None)
-        print(f"[{get_current_time_str()}] 客户端{client_ip}: 资源已清理、在线数: {len(connected_clients)}")
+        print(f"[{get_current_time_str()}] 客户端{client_ip}: 资源已清理，在线数: {len(connected_clients)}")