From cfc6d5f60f3a17095b244015f073f0bfb1ff123c Mon Sep 17 00:00:00 2001
From: Benjamin Braatz <benjamin.braatz@graph-it.com>
Date: Wed, 14 Jul 2021 14:38:45 +0200
Subject: [PATCH] Allow original sender in states and handle lost connection.

---
 controlpi_plugins/graph.py | 39 +++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/controlpi_plugins/graph.py b/controlpi_plugins/graph.py
index f32f4cc..6f38d22 100644
--- a/controlpi_plugins/graph.py
+++ b/controlpi_plugins/graph.py
@@ -30,12 +30,31 @@ class Graph(BasePlugin):
                                        json.dumps(self._states)])
             await self._call('setze', [comessage_guid, 'comessage_ready',
                                        True])
-        if ('state' in message and
-                message['sender'] in self.conf['states']):
-            self._states[message['sender']] = message['state']
+        if 'state' in message:
+            if message['sender'] in self.conf['states']:
+                self._states[message['sender']] = message['state']
+            elif ('original sender' in message and
+                    message['original sender'] in self.conf['states']):
+                self._states[message['original sender']] = message['state']
+
+    async def _open_connection(self):
+        (self._reader, self._writer) = await asyncio.open_connection(
+                self._host, self._port, ssl=self._ssl_ctx)
+        self._writer.transport.set_write_buffer_limits(0)
+        # Read banner:
+        banner_size = await self._reader.readexactly(4)
+        banner_size = struct.unpack('<i', banner_size)[0]
+        banner_message = await self._reader.readexactly(banner_size)
+        # Get node instance:
+        self._call_id = 0
+        self._coroot_guid = await self._call('attributsknoten',
+                                             ['coroot_name',
+                                              self.conf['name']])
 
     async def _call(self, method, params):
         await self._lock.acquire()
+        while self._writer.is_closing():
+            await self._open_connection()
         self._call_id += 1
         request = {'jsonrpc': '2.0', 'method': method,
                    'params': params, 'id': self._call_id}
@@ -70,7 +89,6 @@ class Graph(BasePlugin):
                                     f"'{self.conf['crt']}'.")
         self._ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
         self._ssl_ctx.load_cert_chain(self.conf['crt'])
-        self._call_id = 0
         sends = []
         receives = []
         receives.append(MessageTemplate({'target': {'const': self.name},
@@ -84,15 +102,6 @@ class Graph(BasePlugin):
                           sends, receives, self._receive)
 
     async def run(self) -> None:
-        """Open connection and get node instance for name."""
+        """Open connection and get coroot instance for name."""
         self._lock = asyncio.Lock()
-        (self._reader, self._writer) = await asyncio.open_connection(
-                self._host, self._port, ssl=self._ssl_ctx)
-        # Read banner:
-        banner_size = await self._reader.readexactly(4)
-        banner_size = struct.unpack('<i', banner_size)[0]
-        banner_message = await self._reader.readexactly(banner_size)
-        # Get node instance:
-        self._coroot_guid = await self._call('attributsknoten',
-                                             ['coroot_name',
-                                              self.conf['name']])
+        self._open_connection()
-- 
2.34.1