226 lines
8.1 KiB
Python
226 lines
8.1 KiB
Python
# ///
|
|
# node_manager.py
|
|
# 描述：节点管理服务，负责与WXAUTO-HTTP-API节点通信
|
|
# 作者：AI Generated
|
|
# 创建日期：2026-04-05
|
|
# ///
|
|
|
|
import asyncio
|
|
from typing import Dict, List, Optional, Any
|
|
import httpx
|
|
|
|
from services.log_service import log_service
|
|
from database_service import get_db_service
|
|
|
|
|
|
class NodeStatus:
|
|
ACTIVE = "active"
|
|
INACTIVE = "inactive"
|
|
ERROR = "error"
|
|
|
|
|
|
class Node:
|
|
def __init__(self, node_id: str, name: str, api_url: str, api_key: str, description: str = "", group: str = "default"):
|
|
self.node_id = node_id
|
|
self.name = name
|
|
self.api_url = api_url.rstrip("/")
|
|
self.api_key = api_key
|
|
self.description = description
|
|
self.group = group
|
|
self.enabled = True
|
|
self.status = NodeStatus.INACTIVE
|
|
self.wechat_status = "online"
|
|
self.is_healthy = False
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
"node_id": self.node_id,
|
|
"name": self.name,
|
|
"api_url": self.api_url,
|
|
"enabled": self.enabled,
|
|
"description": self.description,
|
|
"group": self.group,
|
|
"status": self.status,
|
|
"wechat_status": self.wechat_status,
|
|
"is_healthy": self.is_healthy
|
|
}
|
|
|
|
|
|
class NodeManager:
|
|
def __init__(self):
|
|
self.nodes: Dict[str, Node] = {}
|
|
|
|
async def add_node(self, node_id: str, name: str, api_url: str, api_key: str, description: str = "", group: str = "default") -> bool:
|
|
if node_id in self.nodes:
|
|
return False
|
|
self.nodes[node_id] = Node(node_id, name, api_url, api_key, description, group)
|
|
await self.check_node_status(node_id)
|
|
return True
|
|
|
|
async def remove_node(self, node_id: str) -> bool:
|
|
if node_id not in self.nodes:
|
|
return False
|
|
del self.nodes[node_id]
|
|
return True
|
|
|
|
async def get_node(self, node_id: str) -> Optional[Node]:
|
|
return self.nodes.get(node_id)
|
|
|
|
async def get_all_nodes(self) -> List[Node]:
|
|
return list(self.nodes.values())
|
|
|
|
async def get_enabled_nodes(self) -> List[Node]:
|
|
return [n for n in self.nodes.values() if n.enabled]
|
|
|
|
async def get_nodes_by_group(self, group: str) -> List[Node]:
|
|
return [n for n in self.nodes.values() if n.group == group]
|
|
|
|
def to_config(self) -> List[Dict[str, Any]]:
|
|
return [{
|
|
"node_id": n.node_id,
|
|
"name": n.name,
|
|
"api_url": n.api_url,
|
|
"api_key": n.api_key,
|
|
"enabled": n.enabled,
|
|
"description": n.description,
|
|
"group": n.group
|
|
} for n in self.nodes.values()]
|
|
|
|
async def update_node(self, node_id: str, **kwargs) -> bool:
|
|
node = self.nodes.get(node_id)
|
|
if not node:
|
|
return False
|
|
for key, value in kwargs.items():
|
|
if hasattr(node, key):
|
|
setattr(node, key, value)
|
|
if "api_url" in kwargs or "api_key" in kwargs:
|
|
await self.check_node_status(node_id)
|
|
return True
|
|
|
|
async def check_node_status(self, node_id: str) -> str:
|
|
node = self.nodes.get(node_id)
|
|
if not node or not node.enabled:
|
|
return NodeStatus.INACTIVE
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=5.0) as client:
|
|
response = await client.get(
|
|
f"{node.api_url}/health",
|
|
headers={"X-API-Key": node.api_key}
|
|
)
|
|
if response.status_code == 200:
|
|
node.status = NodeStatus.ACTIVE
|
|
else:
|
|
node.status = NodeStatus.ERROR
|
|
except Exception:
|
|
node.status = NodeStatus.ERROR
|
|
return node.status
|
|
|
|
async def is_node_healthy(self, node_id: str) -> tuple[bool, str]:
|
|
node = self.nodes.get(node_id)
|
|
if not node or not node.enabled:
|
|
return False, "Node not found or disabled"
|
|
|
|
api_status = await self.check_node_status(node_id)
|
|
if api_status != NodeStatus.ACTIVE:
|
|
return False, f"API status: {api_status}"
|
|
|
|
return True, "OK"
|
|
|
|
async def call_node_api(self, node_id: str, endpoint: str, method: str = "GET", data: Optional[Dict] = None) -> Dict[str, Any]:
|
|
node = self.nodes.get(node_id)
|
|
if not node or not node.enabled:
|
|
return {"success": False, "error": f"Node '{node_id}' not found or disabled"}
|
|
|
|
url = f"{node.api_url}{endpoint}"
|
|
headers = {"X-API-Key": node.api_key}
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10.0) as client:
|
|
if method.upper() == "GET":
|
|
response = await client.get(url, headers=headers)
|
|
elif method.upper() == "POST":
|
|
response = await client.post(url, headers=headers, json=data)
|
|
else:
|
|
return {"success": False, "error": f"Unsupported method: {method}"}
|
|
|
|
if response.status_code == 200:
|
|
return {"success": True, "data": response.json()}
|
|
else:
|
|
return {"success": False, "error": f"HTTP {response.status_code}", "url": url}
|
|
except Exception as e:
|
|
return {"success": False, "error": str(e), "url": url}
|
|
|
|
async def send_message(self, node_id: str, who: str, msg: str, msg_type: str = "text") -> Dict[str, Any]:
|
|
node = self.nodes.get(node_id)
|
|
await self.call_node_api(node_id, "/api/wechat/initialize", method="POST")
|
|
result = await self.call_node_api(
|
|
node_id,
|
|
"/api/message/send",
|
|
method="POST",
|
|
data={"receiver": who, "message": msg}
|
|
)
|
|
if node:
|
|
node.wechat_status = "online" if result.get("success") else "offline"
|
|
db_service = get_db_service()
|
|
if db_service:
|
|
db_service.update_node(node_id, wechat_status=node.wechat_status)
|
|
return result
|
|
|
|
async def send_message_to_group(self, node_id: str, group_name: str, msg: str, msg_type: str = "text") -> Dict[str, Any]:
|
|
node = self.nodes.get(node_id)
|
|
await self.call_node_api(node_id, "/api/wechat/initialize", method="POST")
|
|
result = await self.call_node_api(
|
|
node_id,
|
|
"/api/message/send",
|
|
method="POST",
|
|
data={"receiver": group_name, "message": msg}
|
|
)
|
|
if node:
|
|
node.wechat_status = "online" if result.get("success") else "offline"
|
|
db_service = get_db_service()
|
|
if db_service:
|
|
db_service.update_node(node_id, wechat_status=node.wechat_status)
|
|
return result
|
|
|
|
async def get_friends_list(self, node_id: str) -> Dict[str, Any]:
|
|
return await self.call_node_api(node_id, "/api/friend/list", method="GET")
|
|
|
|
async def get_groups_list(self, node_id: str) -> Dict[str, Any]:
|
|
return await self.call_node_api(node_id, "/api/group/list", method="GET")
|
|
|
|
async def get_messages(self, node_id: str, chat_name: str, limit: int = 20) -> Dict[str, Any]:
|
|
return await self.call_node_api(
|
|
node_id,
|
|
"/api/message/get-history",
|
|
method="POST",
|
|
data={"chat_name": chat_name, "save_pic": False, "save_video": False, "save_file": False, "save_voice": False}
|
|
)
|
|
|
|
async def reset_wechat_status(self, node_id: str) -> Dict[str, Any]:
|
|
node = self.nodes.get(node_id)
|
|
if not node:
|
|
return {"success": False, "error": f"Node '{node_id}' not found"}
|
|
|
|
await self.call_node_api(node_id, "/api/wechat/initialize", method="POST")
|
|
result = await self.call_node_api(
|
|
node_id,
|
|
"/api/message/send",
|
|
method="POST",
|
|
data={"receiver": "文件传输助手", "message": "WXAuto Center 状态检测消息"}
|
|
)
|
|
|
|
if result.get("success"):
|
|
node.wechat_status = "online"
|
|
else:
|
|
node.wechat_status = "offline"
|
|
|
|
db_service = get_db_service()
|
|
if db_service:
|
|
db_service.update_node(node_id, wechat_status=node.wechat_status)
|
|
|
|
return result
|
|
|
|
|
|
node_manager = NodeManager()
|