Uni-Lab-OS/unilabos/app/mq.py

import json
import time
import traceback
import uuid

import paho.mqtt.client as mqtt
import ssl
import base64
import hmac
from hashlib import sha1
import tempfile
import os

from unilabos.config.config import MQConfig
from unilabos.app.controler import devices, job_add
from unilabos.app.model import JobAddReq
from unilabos.utils import logger
from unilabos.utils.type_check import TypeEncoder

from paho.mqtt.enums import CallbackAPIVersion


class MQTTClient:
    mqtt_disable = True

    def __init__(self):
        self.mqtt_disable = not MQConfig.lab_id
        self.client_id = f"{MQConfig.group_id}@@@{MQConfig.lab_id}{uuid.uuid4()}"
        logger.info("[MQTT] Client_id: " + self.client_id)
        self.client = mqtt.Client(CallbackAPIVersion.VERSION2, client_id=self.client_id, protocol=mqtt.MQTTv5)
        self._setup_callbacks()

    def _setup_callbacks(self):
        self.client.on_log = self._on_log
        self.client.on_connect = self._on_connect
        self.client.on_message = self._on_message
        self.client.on_disconnect = self._on_disconnect

    def _on_log(self, client, userdata, level, buf):
        # logger.info(f"[MQTT] log: {buf}")
        pass

    def _on_connect(self, client, userdata, flags, rc, properties=None):
        logger.info("[MQTT] Connected with result code " + str(rc))
        client.subscribe(f"labs/{MQConfig.lab_id}/job/start/", 0)
        isok, data = devices()
        if not isok:
            logger.error("[MQTT] on_connect ErrorHostNotInit")
            return

    def _on_message(self, client, userdata, msg) -> None:
        logger.info("[MQTT] on_message<<<< " + msg.topic + " " + str(msg.payload))
        try:
            payload_str = msg.payload.decode("utf-8")
            payload_json = json.loads(payload_str)
            logger.debug(f"Topic: {msg.topic}")
            logger.debug("Payload:", json.dumps(payload_json, indent=2, ensure_ascii=False))
            if msg.topic == f"labs/{MQConfig.lab_id}/job/start/":
                logger.debug("job_add", type(payload_json), payload_json)
                if "data" not in payload_json:
                    payload_json["data"] = {}
                if "action" in payload_json:
                    payload_json["data"]["action"] = payload_json.pop("action")
                if "action_kwargs" in payload_json:
                    payload_json["data"]["action_kwargs"] = payload_json.pop("action_kwargs")
                job_req = JobAddReq.model_validate(payload_json)
                data = job_add(job_req)
                return

        except json.JSONDecodeError as e:
            logger.error(f"[MQTT] JSON 解析错误: {e}")
            logger.error(f"[MQTT] Raw message: {msg.payload}")
            logger.error(traceback.format_exc())
        except Exception as e:
            logger.error(f"[MQTT] 处理消息时出错: {e}")
            logger.error(traceback.format_exc())

    def _on_disconnect(self, client, userdata, rc, reasonCode=None, properties=None):
        if rc != 0:
            logger.error(f"[MQTT] Unexpected disconnection {rc}")

    def _setup_ssl_context(self):
        temp_files = []
        try:
            with tempfile.NamedTemporaryFile(mode="w", delete=False) as ca_temp:
                ca_temp.write(MQConfig.ca_content)
                temp_files.append(ca_temp.name)

            with tempfile.NamedTemporaryFile(mode="w", delete=False) as cert_temp:
                cert_temp.write(MQConfig.cert_content)
                temp_files.append(cert_temp.name)

            with tempfile.NamedTemporaryFile(mode="w", delete=False) as key_temp:
                key_temp.write(MQConfig.key_content)
                temp_files.append(key_temp.name)

            context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
            context.load_verify_locations(cafile=temp_files[0])
            context.load_cert_chain(certfile=temp_files[1], keyfile=temp_files[2])
            self.client.tls_set_context(context)
        finally:
            for temp_file in temp_files:
                try:
                    os.unlink(temp_file)
                except Exception as e:
                    pass

    def start(self):
        if self.mqtt_disable:
            logger.warning("MQTT is disabled, skipping connection.")
            return
        userName = f"Signature|{MQConfig.access_key}|{MQConfig.instance_id}"
        password = base64.b64encode(
            hmac.new(MQConfig.secret_key.encode(), self.client_id.encode(), sha1).digest()
        ).decode()

        self.client.username_pw_set(userName, password)
        self._setup_ssl_context()

        # 创建连接线程
        def connect_thread_func():
            try:
                self.client.connect(MQConfig.broker_url, MQConfig.port, 60)
                self.client.loop_start()

                # 添加连接超时检测
                max_attempts = 5
                attempt = 0
                while not self.client.is_connected() and attempt < max_attempts:
                    logger.info(
                        f"[MQTT] 正在连接到 {MQConfig.broker_url}:{MQConfig.port}，尝试 {attempt+1}/{max_attempts}"
                    )
                    time.sleep(3)
                    attempt += 1

                if self.client.is_connected():
                    logger.info(f"[MQTT] 已成功连接到 {MQConfig.broker_url}:{MQConfig.port}")
                else:
                    logger.error(f"[MQTT] 连接超时，可能是账号密码错误或网络问题")
                    self.client.loop_stop()
            except Exception as e:
                logger.error(f"[MQTT] 连接失败: {str(e)}")

        connect_thread_func()
        # connect_thread = threading.Thread(target=connect_thread_func)
        # connect_thread.daemon = True
        # connect_thread.start()

    def stop(self):
        if self.mqtt_disable:
            return
        self.client.disconnect()
        self.client.loop_stop()

    def publish_device_status(self, device_status: dict, device_id, property_name):
        # status = device_status.get(device_id, {})
        if self.mqtt_disable:
            return
        status = {"data": device_status.get(device_id, {}), "device_id": device_id}
        address = f"labs/{MQConfig.lab_id}/devices/"
        self.client.publish(address, json.dumps(status), qos=2)
        logger.critical(f"Device status published: address: {address}, {status}")

    def publish_job_status(self, feedback_data: dict, job_id: str, status: str):
        if self.mqtt_disable:
            return
        jobdata = {"job_id": job_id, "data": feedback_data, "status": status}
        self.client.publish(f"labs/{MQConfig.lab_id}/job/list/", json.dumps(jobdata), qos=2)

    def publish_registry(self, device_id: str, device_info: dict):
        if self.mqtt_disable:
            return
        address = f"labs/{MQConfig.lab_id}/registry/"
        registry_data = json.dumps({device_id: device_info}, ensure_ascii=False, cls=TypeEncoder)
        self.client.publish(address, registry_data, qos=2)
        logger.debug(f"Registry data published: address: {address}, {registry_data}")

    def publish_actions(self, action_id: str, action_info: dict):
        if self.mqtt_disable:
            return
        address = f"labs/{MQConfig.lab_id}/actions/"
        self.client.publish(address, json.dumps(action_info), qos=2)
        logger.debug(f"Action data published: address: {address}, {action_id}, {action_info}")


mqtt_client = MQTTClient()

if __name__ == "__main__":
    mqtt_client.start()