pigway-frame-flow/pigway_frame_flow/extensions/mqtt.py

"""Opt-in MQTT ingress shares the display queue, scopes and idempotency."""

import hashlib
import json

from fastapi import Depends
from pydantic import Field

from pigway_frame_flow.core.auth import require
from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.models import DisplayRequest, StrictModel
from pigway_frame_flow.settings import secret
from pigway_frame_flow.storage.database import encode


class MQTTConfig(StrictModel):
    enabled: bool = False
    host: str = Field(default="", pattern=r"^[a-zA-Z0-9.:-]*$", max_length=253)
    port: int = Field(default=8883, ge=1, le=65535)
    tls: bool = True
    allow_insecure_lan: bool = False
    topic: str = Field(default="frame-flow/display", pattern=r"^[a-zA-Z0-9_/-]{1,120}$")
    username: str = Field(default="", max_length=120)


class MQTTBridge:
    def __init__(self, service):
        self.service = service
        self.client = None
        self.connected = False

    def config(self):
        row = self.service.db.one("SELECT payload FROM app_config WHERE key='mqtt'")
        return MQTTConfig.model_validate_json(row["payload"]) if row else MQTTConfig()

    def accept(self, payload, retained=False):
        if retained:
            raise FlowError(E.INVALID_REQUEST, "Retained display commands are rejected")
        if len(payload) > 65536:
            raise FlowError(E.INVALID_REQUEST, "MQTT command exceeds size limit")
        token = secret("FRAME_FLOW_MQTT_TOKEN")
        row = (
            self.service.db.one(
                "SELECT * FROM clients WHERE token_hash=? AND enabled=1",
                (hashlib.sha256(token.encode()).hexdigest(),),
            )
            if token
            else None
        )
        if not row:
            raise FlowError(
                E.UNAUTHORIZED, "Configure an enabled Frame Flow MQTT client token", 401
            )
        scopes = json.loads(row["permissions"])
        if "admin" not in scopes and "display:write" not in scopes:
            raise FlowError(E.FORBIDDEN, "MQTT client lacks display:write", 403)
        self.service.settings.additional_secrets.append(token)
        data = json.loads(payload)
        request_id = data.pop("request_id", "")
        if not isinstance(request_id, str) or not request_id or len(request_id) > 120:
            raise FlowError(
                E.INVALID_REQUEST, "MQTT commands require a stable request_id"
            )
        request = DisplayRequest.model_validate(data)
        self.service.validate(request)
        id = self.service.submit(
            "display",
            request.model_dump(mode="json"),
            request.device_id,
            row["client_id"],
            self.service.settings.sanitize(request_id),
            "mqtt-" + hashlib.sha256(request_id.encode()).hexdigest(),
        )
        return {"operation_id": id, "status": "accepted"}

    def start(self):
        config = self.config()
        if not config.enabled:
            return
        if not config.host or (not config.tls and not config.allow_insecure_lan):
            raise FlowError(
                E.CONFIG_REQUIRED,
                "MQTT requires host and TLS or explicit LAN opt-in",
                409,
            )
        import paho.mqtt.client as mqtt

        self.client = mqtt.Client(
            mqtt.CallbackAPIVersion.VERSION2,
            client_id="frame-flow-"
            + hashlib.sha256(str(self.service.settings.data_dir).encode()).hexdigest()[
                :12
            ],
        )
        password = secret("FRAME_FLOW_MQTT_PASSWORD")
        self.service.settings.additional_secrets.append(password)
        if config.username:
            self.client.username_pw_set(config.username, password)
        if config.tls:
            self.client.tls_set()
        self.client.reconnect_delay_set(min_delay=5, max_delay=120)

        def connected(client, userdata, flags, reason_code, properties):
            self.connected = not reason_code.is_failure
            if self.connected:
                client.subscribe(config.topic, qos=1)

        def disconnected(client, userdata, flags, reason_code, properties):
            self.connected = False

        def message(client, userdata, message):
            try:
                response = self.accept(message.payload, message.retain)
            except FlowError as exc:
                response = {"error": {"code": exc.code, "message": exc.message}}
            except Exception:
                response = {
                    "error": {
                        "code": E.INVALID_REQUEST,
                        "message": "Invalid MQTT display request",
                    }
                }
            client.publish(
                config.topic + "/result", encode(response), qos=1, retain=False
            )

        self.client.on_connect = connected
        self.client.on_disconnect = disconnected
        self.client.on_message = message
        self.client.connect_async(config.host, config.port, keepalive=60)
        self.client.loop_start()

    def stop(self):
        if self.client:
            self.client.disconnect()
            self.client.loop_stop()
            self.client = None
        self.connected = False


def install_mqtt(app, service):
    service.mqtt = MQTTBridge(service)
    app.state.mqtt = service.mqtt

    @app.get(
        "/api/v1/integrations/mqtt",
        dependencies=[Depends(require("admin"))],
        tags=["MQTT"],
    )
    def state():
        return {
            "config": service.mqtt.config().model_dump(),
            "connected": service.mqtt.connected,
        }

    @app.put(
        "/api/v1/integrations/mqtt",
        dependencies=[Depends(require("admin"))],
        tags=["MQTT"],
    )
    def configure(payload: MQTTConfig):
        if payload.enabled and (
            not payload.host or not payload.tls and not payload.allow_insecure_lan
        ):
            raise FlowError(
                E.INVALID_REQUEST, "MQTT needs a host and TLS or explicit LAN opt-in"
            )
        service.mqtt.stop()
        service.db.execute(
            "INSERT OR REPLACE INTO app_config VALUES('mqtt',?)",
            (payload.model_dump_json(),),
        )
        if getattr(app.state, "worker_enabled", False):
            service.mqtt.start()
        return state()