From a9f9a67184d8208a82f41ab550b475cb67f02e90 Mon Sep 17 00:00:00 2001
From: mokaddem <sami.mokaddem@circl.lu>
Date: Fri, 21 Jun 2019 15:55:03 +0200
Subject: [PATCH] chg: [diagnostic] Added support of multiple subscribers - WiP

---
 diagnostic.py | 43 ++++++++++++++++++++++++++++---------------
 1 file changed, 28 insertions(+), 15 deletions(-)

diff --git a/diagnostic.py b/diagnostic.py
index 8d0faeb..a51e8a5 100755
--- a/diagnostic.py
+++ b/diagnostic.py
@@ -172,22 +172,35 @@ def check_zmq(spinner):
     timeout = 15
     context = zmq.Context()
     socket = context.socket(zmq.SUB)
-    socket.connect(configuration_file.get('RedisGlobal', 'zmq_url'))
-    socket.setsockopt_string(zmq.SUBSCRIBE, '')
-    poller = zmq.Poller()
+    misp_instances = json.loads(cfg.get('RedisGlobal', 'misp_instances'))
+    instances_status = {}
+    for misp_instance in misp_instances:
+        socket.connect(misp_instance.get('zmq'))
+        socket.setsockopt_string(zmq.SUBSCRIBE, '')
+        poller = zmq.Poller()
 
-    start_time = time.time()
-    poller.register(socket, zmq.POLLIN)
-    for t in range(1, timeout+1):
-        socks = dict(poller.poll(timeout=1*1000))
-        if len(socks) > 0:
-            if socket in socks and socks[socket] == zmq.POLLIN:
-                rcv_string = socket.recv()
-                if rcv_string.startswith(b'misp_json'):
-                    return (True, '')
-        else:
-            pass
-            spinner.text = f'checking zmq - elapsed time: {int(time.time() - start_time)}s'
+        start_time = time.time()
+        poller.register(socket, zmq.POLLIN)
+        for t in range(1, timeout+1):
+            socks = dict(poller.poll(timeout=1*1000))
+            if len(socks) > 0:
+                if socket in socks and socks[socket] == zmq.POLLIN:
+                    rcv_string = socket.recv()
+                    if rcv_string.startswith(b'misp_json'):
+                        instances_status[misp_instance.get('name')] = True
+            else:
+                pass
+                spinner.text = f'checking zmq of {misp_instance.get('name')} - elapsed time: {int(time.time() - start_time)}s'
+        instances_status[misp_instance.get('name')] = False
+
+    results = [s for n, s in instances_status.items()]
+    if all(results):
+        return (True, '')
+    elif any(results):
+        return_text = 'Some connection to ZMQ streams failed.\n'
+        for name, status:
+            return_text += f'\t➥ {name}: {'success' if status else 'failed'}\n'
+        return (True, return_text)
     else:
         return (False, '''Can\'t connect to the ZMQ stream.
 \t➥ Make sure the MISP ZMQ is running: `/servers/serverSettings/diagnostics`