Make ws client conform to communicator

This commit is contained in:
Nick Mowen 2022-11-22 13:29:28 -07:00
parent c4a2ac3726
commit 0ca8219fbd

View File

@ -1,22 +1,24 @@
"""Websocket communicator.""" """Websocket communicator."""
import json
import logging
import threading
from wsgiref.simple_server import make_server
from ws4py.server.wsgirefserver import ( from ws4py.server.wsgirefserver import (
WebSocketWSGIHandler, WebSocketWSGIHandler,
WebSocketWSGIRequestHandler, WebSocketWSGIRequestHandler,
WSGIServer, WSGIServer,
) )
from ws4py.server.wsgiutils import WebSocketWSGIApplication
from ws4py.websocket import WebSocket
from frigate.communication.dispatcher import Communicator from frigate.communication.dispatcher import Communicator
class MqttSocketRelay(Communicator): logger = logging.getLogger(__name__)
def __init__(self, topic_prefix: str):
self.topic_prefix = topic_prefix
def start(self):
class MqttWebSocket(WebSocket):
topic_prefix = self.topic_prefix
mqtt_client = self.mqtt_client
class WebSocketHandler(WebSocket):
def received_message(self, message): def received_message(self, message):
try: try:
json_message = json.loads(message.data.decode("utf-8")) json_message = json.loads(message.data.decode("utf-8"))
@ -32,12 +34,21 @@ class MqttSocketRelay(Communicator):
logger.debug( logger.debug(
f"Publishing mqtt message from websockets at {json_message['topic']}." f"Publishing mqtt message from websockets at {json_message['topic']}."
) )
self.mqtt_client.publish( self.publish(
json_message["topic"], json_message["topic"],
json_message["payload"], json_message["payload"],
retain=json_message["retain"], retain=json_message["retain"],
) )
class WebSocketClient(Communicator):
"""Frigate wrapper for ws client."""
def __init__(self) -> None:
pass
def start(self):
# start a websocket server on 5002 # start a websocket server on 5002
WebSocketWSGIHandler.http_version = "1.1" WebSocketWSGIHandler.http_version = "1.1"
self.websocket_server = make_server( self.websocket_server = make_server(
@ -45,36 +56,29 @@ class MqttSocketRelay(Communicator):
5002, 5002,
server_class=WSGIServer, server_class=WSGIServer,
handler_class=WebSocketWSGIRequestHandler, handler_class=WebSocketWSGIRequestHandler,
app=WebSocketWSGIApplication(handler_cls=MqttWebSocket), app=WebSocketWSGIApplication(handler_cls=WebSocketHandler),
) )
self.websocket_server.initialize_websockets_manager() self.websocket_server.initialize_websockets_manager()
self.websocket_thread = threading.Thread( self.websocket_thread = threading.Thread(
target=self.websocket_server.serve_forever target=self.websocket_server.serve_forever
) )
self.websocket_thread.start()
def send(client, userdata, message): def publish(self, topic: str, payload: str) -> None:
"""Sends mqtt messages to clients."""
try: try:
logger.debug(f"Received mqtt message on {message.topic}.")
ws_message = json.dumps( ws_message = json.dumps(
{ {
"topic": message.topic.replace(f"{self.topic_prefix}/", ""), "topic": topic,
"payload": message.payload.decode(), "payload": payload,
} }
) )
except Exception as e: except Exception as e:
# if the payload can't be decoded don't relay to clients # if the payload can't be decoded don't relay to clients
logger.debug( logger.debug(f"payload for {topic} wasn't text. Skipping...")
f"MQTT payload for {message.topic} wasn't text. Skipping..."
)
return return
self.websocket_server.manager.broadcast(ws_message) self.websocket_server.manager.broadcast(ws_message)
self.mqtt_client.add_topic_callback(f"{self.topic_prefix}/#", send)
self.websocket_thread.start()
def stop(self): def stop(self):
self.websocket_server.manager.close_all() self.websocket_server.manager.close_all()
self.websocket_server.manager.stop() self.websocket_server.manager.stop()