250 lines
8.0 KiB
Python
250 lines
8.0 KiB
Python
"""
|
|
设备WebSocket路由
|
|
接收Android Agent APP的连接和命令
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
from typing import Dict, Optional
|
|
import json
|
|
import asyncio
|
|
import logging
|
|
from datetime import datetime
|
|
|
|
router = APIRouter(tags=["WebSocket"])
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DeviceConnectionManager:
|
|
"""
|
|
设备连接管理器
|
|
管理所有连接的Android设备
|
|
"""
|
|
|
|
def __init__(self):
|
|
# device_id -> WebSocket连接
|
|
self.active_connections: Dict[str, WebSocket] = {}
|
|
# device_id -> 设备信息
|
|
self.device_info: Dict[str, dict] = {}
|
|
# project_id -> [device_id, ...]
|
|
self.project_devices: Dict[str, list] = {}
|
|
|
|
async def connect(self, device_id: str, websocket: WebSocket):
|
|
"""接受新连接"""
|
|
await websocket.accept()
|
|
self.active_connections[device_id] = websocket
|
|
logger.info(f"设备连接: {device_id}")
|
|
|
|
def disconnect(self, device_id: str):
|
|
"""断开连接"""
|
|
if device_id in self.active_connections:
|
|
del self.active_connections[device_id]
|
|
if device_id in self.device_info:
|
|
project_id = self.device_info[device_id].get("project_id")
|
|
if project_id and project_id in self.project_devices:
|
|
if device_id in self.project_devices[project_id]:
|
|
self.project_devices[project_id].remove(device_id)
|
|
del self.device_info[device_id]
|
|
logger.info(f"设备断开: {device_id}")
|
|
|
|
def register_device(self, device_id: str, info: dict):
|
|
"""注册设备信息"""
|
|
self.device_info[device_id] = {
|
|
**info,
|
|
"connected_at": datetime.now().isoformat(),
|
|
"last_heartbeat": datetime.now().isoformat()
|
|
}
|
|
|
|
# 添加到项目
|
|
project_id = info.get("project_id")
|
|
if project_id:
|
|
if project_id not in self.project_devices:
|
|
self.project_devices[project_id] = []
|
|
if device_id not in self.project_devices[project_id]:
|
|
self.project_devices[project_id].append(device_id)
|
|
|
|
logger.info(f"设备注册: {device_id}, 项目: {project_id}")
|
|
|
|
def update_heartbeat(self, device_id: str):
|
|
"""更新心跳时间"""
|
|
if device_id in self.device_info:
|
|
self.device_info[device_id]["last_heartbeat"] = datetime.now().isoformat()
|
|
|
|
async def send_to_device(self, device_id: str, message: dict) -> bool:
|
|
"""发送消息到指定设备"""
|
|
if device_id in self.active_connections:
|
|
try:
|
|
await self.active_connections[device_id].send_json(message)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"发送失败: {device_id}, {e}")
|
|
return False
|
|
return False
|
|
|
|
async def broadcast_to_project(self, project_id: str, message: dict):
|
|
"""广播消息到项目下所有设备"""
|
|
if project_id in self.project_devices:
|
|
for device_id in self.project_devices[project_id]:
|
|
await self.send_to_device(device_id, message)
|
|
|
|
def get_device(self, device_id: str) -> Optional[dict]:
|
|
"""获取设备信息"""
|
|
return self.device_info.get(device_id)
|
|
|
|
def get_project_devices(self, project_id: str) -> list:
|
|
"""获取项目下所有设备"""
|
|
device_ids = self.project_devices.get(project_id, [])
|
|
return [self.device_info[did] for did in device_ids if did in self.device_info]
|
|
|
|
def get_all_devices(self) -> list:
|
|
"""获取所有设备"""
|
|
return list(self.device_info.values())
|
|
|
|
def is_connected(self, device_id: str) -> bool:
|
|
"""检查设备是否在线"""
|
|
return device_id in self.active_connections
|
|
|
|
|
|
# 全局连接管理器
|
|
manager = DeviceConnectionManager()
|
|
|
|
|
|
@router.websocket("/ws/device/{device_id}")
|
|
async def websocket_endpoint(websocket: WebSocket, device_id: str):
|
|
"""
|
|
设备WebSocket连接端点
|
|
|
|
消息类型:
|
|
- register: 设备注册
|
|
- heartbeat: 心跳
|
|
- result: 命令执行结果
|
|
"""
|
|
await manager.connect(device_id, websocket)
|
|
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
msg_type = data.get("type")
|
|
|
|
if msg_type == "register":
|
|
# 设备注册
|
|
manager.register_device(device_id, data)
|
|
await websocket.send_json({
|
|
"type": "registered",
|
|
"device_id": device_id,
|
|
"message": "注册成功"
|
|
})
|
|
|
|
elif msg_type == "heartbeat":
|
|
# 心跳
|
|
manager.update_heartbeat(device_id)
|
|
await websocket.send_json({
|
|
"type": "pong",
|
|
"timestamp": datetime.now().timestamp() * 1000
|
|
})
|
|
|
|
elif msg_type == "result":
|
|
# 命令执行结果
|
|
command_id = data.get("command_id")
|
|
logger.info(f"命令结果: {command_id}, 成功: {data.get('success')}")
|
|
# 这里可以存储结果或通知其他系统
|
|
|
|
else:
|
|
logger.warning(f"未知消息类型: {msg_type}")
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(device_id)
|
|
except Exception as e:
|
|
logger.error(f"WebSocket错误: {device_id}, {e}")
|
|
manager.disconnect(device_id)
|
|
|
|
|
|
# ========== REST API接口 ==========
|
|
|
|
@router.get("/api/v3/devices")
|
|
async def list_devices():
|
|
"""获取所有在线设备"""
|
|
return {
|
|
"success": True,
|
|
"devices": manager.get_all_devices(),
|
|
"total": len(manager.get_all_devices())
|
|
}
|
|
|
|
|
|
@router.get("/api/v3/devices/{device_id}")
|
|
async def get_device(device_id: str):
|
|
"""获取指定设备信息"""
|
|
device = manager.get_device(device_id)
|
|
if device:
|
|
return {
|
|
"success": True,
|
|
"device": device,
|
|
"online": manager.is_connected(device_id)
|
|
}
|
|
return {"success": False, "message": "设备不存在"}
|
|
|
|
|
|
@router.get("/api/v3/projects/{project_id}/devices")
|
|
async def get_project_devices(project_id: str):
|
|
"""获取项目下所有设备"""
|
|
devices = manager.get_project_devices(project_id)
|
|
return {
|
|
"success": True,
|
|
"project_id": project_id,
|
|
"devices": devices,
|
|
"total": len(devices)
|
|
}
|
|
|
|
|
|
@router.post("/api/v3/devices/{device_id}/execute")
|
|
async def execute_command(device_id: str, command: dict):
|
|
"""
|
|
向设备发送执行命令
|
|
|
|
请求体:
|
|
{
|
|
"action": "open_app",
|
|
"params": {"package": "com.tencent.mm"}
|
|
}
|
|
"""
|
|
if not manager.is_connected(device_id):
|
|
return {"success": False, "message": "设备离线"}
|
|
|
|
command_id = f"cmd_{datetime.now().timestamp()}"
|
|
message = {
|
|
"type": "execute",
|
|
"command_id": command_id,
|
|
"action": command.get("action"),
|
|
"params": command.get("params", {})
|
|
}
|
|
|
|
success = await manager.send_to_device(device_id, message)
|
|
|
|
return {
|
|
"success": success,
|
|
"command_id": command_id,
|
|
"message": "命令已发送" if success else "发送失败"
|
|
}
|
|
|
|
|
|
@router.post("/api/v3/projects/{project_id}/broadcast")
|
|
async def broadcast_to_project(project_id: str, command: dict):
|
|
"""向项目下所有设备广播命令"""
|
|
command_id = f"cmd_{datetime.now().timestamp()}"
|
|
message = {
|
|
"type": "execute",
|
|
"command_id": command_id,
|
|
"action": command.get("action"),
|
|
"params": command.get("params", {})
|
|
}
|
|
|
|
await manager.broadcast_to_project(project_id, message)
|
|
|
|
devices = manager.get_project_devices(project_id)
|
|
return {
|
|
"success": True,
|
|
"command_id": command_id,
|
|
"devices_count": len(devices),
|
|
"message": f"已广播到 {len(devices)} 台设备"
|
|
}
|