114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
import paho.mqtt.client as mqtt
|
||
import json
|
||
import logging
|
||
import threading
|
||
from typing import Optional
|
||
import uuid
|
||
import threading
|
||
import requests
|
||
from os import getenv
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
MQTT_CLIENT_ID = f"MCPMQTT-{uuid.uuid4().hex[:8]}"
|
||
|
||
def getConfig_url(array):
|
||
url = "http://lzwcai-demp-corp-manager:8086/system/config/getConfig"
|
||
data = [array]
|
||
try:
|
||
response = requests.post(url, json=data, timeout=5)
|
||
response.raise_for_status()
|
||
data = response.json()['data']
|
||
return data[0]['configValue']
|
||
except Exception as e:
|
||
print(f"Error fetching config for {array}: {e}")
|
||
return None
|
||
|
||
class MQTTHandler:
|
||
def __init__(self):
|
||
self.client = mqtt.Client()
|
||
self.tasks = {} # {task_id: task_data}
|
||
self._lock = threading.Lock()
|
||
|
||
# 获取MQTT配置，提供默认值
|
||
self.mqtt_username = getenv('MQTT_USERNAME') or 'lzwc'
|
||
self.mqtt_password = getenv('MQTT_PASSWORD') or 'Lzwc@4187.'
|
||
mqtt_broker_raw = getenv('MQTT_BROKER') or 'tcp://dempmana.lzwcai.com'
|
||
# 移除协议前缀，只保留主机名
|
||
self.mqtt_broker = mqtt_broker_raw.replace('tcp://', '').replace('mqtt://', '')
|
||
mqtt_port_str = getenv('MQTT_PORT')
|
||
self.mqtt_port = int(mqtt_port_str) if mqtt_port_str else 1883
|
||
|
||
# 记录配置信息
|
||
logger.info(f"MQTT配置 - Broker: {self.mqtt_broker}, Port: {self.mqtt_port}, Username: {self.mqtt_username}")
|
||
if not getenv('MQTT_BROKER'):
|
||
logger.warning("MQTT_BROKER环境变量未设置，使用默认值")
|
||
if not mqtt_port_str:
|
||
logger.warning("MQTT_PORT环境变量未设置，使用默认端口1883")
|
||
|
||
self.client.username_pw_set(self.mqtt_username, self.mqtt_password)
|
||
self.client.on_connect = self._on_connect
|
||
self.client.on_message = self._on_message
|
||
self.client.on_disconnect = self._on_disconnect
|
||
try:
|
||
logger.info(f"正在连接 MQTT 代理: {self.mqtt_broker}:{self.mqtt_port}")
|
||
self.client.connect(self.mqtt_broker, self.mqtt_port, 60)
|
||
self.client.loop_start()
|
||
except Exception as e:
|
||
logger.error(f"连接 MQTT 失败: {e}")
|
||
|
||
|
||
def _on_connect(self, client, userdata, flags, rc):
|
||
"""MQTT 连接回调"""
|
||
if rc == 0:
|
||
logger.info("已成功连接到 MQTT 代理服务器")
|
||
else:
|
||
logger.error(f"连接 MQTT 代理服务器失败，返回码: {rc}")
|
||
|
||
def get_status(self, task_id: str = None):
|
||
"""获取任务状态
|
||
Args:
|
||
task_id: 可选参数，指定任务ID。如果为None则返回所有任务状态
|
||
Returns:
|
||
如果指定task_id则返回该任务的状态，否则返回所有任务状态
|
||
"""
|
||
with self._lock:
|
||
if task_id:
|
||
return self.tasks.get(task_id, {}).copy()
|
||
return {k: v.copy() for k, v in self.tasks.items()}
|
||
|
||
def _on_message(self, client, userdata, msg):
|
||
try:
|
||
payload = json.loads(msg.payload.decode())
|
||
logger.info(f"[MQTT] topic={msg.topic} msg: {payload}")
|
||
|
||
# task_id = payload.get('task_id')
|
||
# if task_id and task_id in self.tasks:
|
||
# with self._lock:
|
||
# self.tasks[task_id].update({
|
||
# 'status': payload.get('status', 'UNKNOWN'),
|
||
# 'description': payload.get('description', '')
|
||
# })
|
||
|
||
except Exception as e:
|
||
logger.error(f"[错误] 处理消息失败: {e}")
|
||
|
||
def _on_disconnect(self, client, userdata, rc):
|
||
"""MQTT 断开连接回调"""
|
||
if rc != 0:
|
||
logger.warning("意外断开与 MQTT 代理服务器的连接。")
|
||
try:
|
||
logger.info("尝试重新连接 MQTT...")
|
||
client.reconnect()
|
||
except Exception as e:
|
||
logger.error(f"重新连接 MQTT 失败: {e}")
|
||
|
||
# 单例模式
|
||
_instance: Optional[MQTTHandler] = None
|
||
|
||
def get_mcpmqtt_handler() -> MQTTHandler:
|
||
"""获取单例"""
|
||
global _instance
|
||
if _instance is None:
|
||
_instance = MQTTHandler()
|
||
return _instance
|