From ce9afa67865c77f1ebdedc33863f5987de50b0ab Mon Sep 17 00:00:00 2001
From: Benjamin Braatz <benjamin.braatz@graph-it.com>
Date: Wed, 19 May 2021 10:47:15 +0200
Subject: [PATCH] Allow configuration (name and filters) of websocket.

---
 controlpi_plugins/wsserver.py | 51 +++++++++++++++++++++++++----------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/controlpi_plugins/wsserver.py b/controlpi_plugins/wsserver.py
index 2a79b8c..d2ce1b0 100644
--- a/controlpi_plugins/wsserver.py
+++ b/controlpi_plugins/wsserver.py
@@ -5,7 +5,6 @@
 TODO: documentation, doctests
 TODO: Let Debug web app collapse/expand nested structures
 TODO: Make Debug web app work with nested structures in commands
-TODO: Let clients filter messages received over websocket
 """
 import os
 import sys
@@ -47,25 +46,49 @@ class Connection:
         json_message = json.dumps(message)
         try:
             await self._websocket.send(json_message)
-        except (ConnectionClosedOK, ConnectionClosedError):
+        except ConnectionClosed:
             pass
 
     async def run(self):
         """Listen on websocket and relay messages to bus."""
-        await self._bus.send({'sender': self._name,
-                              'event': 'connection opened',
-                              'address': self._address,
-                              'port': self._port})
+        await self._bus.send(Message(self._name,
+                                     {'event': 'connection opened',
+                                      'address': self._address,
+                                      'port': self._port}))
         try:
             async for json_message in self._websocket:
-                original_message = json.loads(json_message)
-                message = {'sender': self._name}
-                message.update(original_message)
-                await self._bus.send(message)
+                message = json.loads(json_message)
+                if ('command' in message and
+                        message['command'] == 'configure websocket' and
+                        'target' in message and
+                        message['target'] == ''):
+                    await self._bus.send(Message(self._name,
+                                                 {'event':
+                                                  'connection closed'}))
+                    self._bus.unregister(self._name)
+                    self._name = message['name']
+                    sends = []
+                    sends.append(MessageTemplate({'event':
+                                                  {'const':
+                                                   'connection opened'}}))
+                    sends.append(MessageTemplate({'event':
+                                                  {'const':
+                                                   'connection closed'}}))
+                    for template in message['up filter']:
+                        sends.append(template)
+                    self._bus.register(self._name, 'WSServer', sends,
+                                       message['down filter'], self._receive)
+                    await self._bus.send(Message(self._name,
+                                                 {'event':
+                                                  'connection opened',
+                                                  'address': self._address,
+                                                  'port': self._port}))
+                else:
+                    await self._bus.send(Message(self._name, message))
         except ConnectionClosed:
             pass
-        await self._bus.send({'sender': self._name,
-                              'event': 'connection closed'})
+        await self._bus.send(Message(self._name,
+                                     {'event': 'connection closed'}))
         self._bus.unregister(self._name)
 
 
@@ -73,7 +96,7 @@ Response = Optional[Tuple[HTTPStatus, Headers, bytes]]
 
 
 class WSServer(BasePlugin):
-    """Websocket server as ControlPi plugin.
+    """Websocket server plugin.
 
     Run websocket server on host and port given in configuration, serving
     the contents in given web root.
@@ -90,7 +113,7 @@ class WSServer(BasePlugin):
                                              'location': {'type': 'string'}},
                               'required': ['location']}},
                             'additionalProperties': False}}}
-    """Schema for WServer plugin configuration.
+    """Schema for WSServer plugin configuration.
 
     Optional configuration keys:
 
-- 
2.34.1