From 57b4ec4117309f9e9dbf93066930b527a2f9c5d0 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Mon, 3 Oct 2022 18:15:15 +0000
Subject: [PATCH] BB: Refactor agent communication check

Updated CommunicationAnalyzer to use the /api/agents and /api/machines
endpoints to determine whether or not an agent communicated back to the
island.

Resolves PR #2388
---
 .../analyzers/communication_analyzer.py       | 28 ++++++++++++-------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/envs/monkey_zoo/blackbox/analyzers/communication_analyzer.py b/envs/monkey_zoo/blackbox/analyzers/communication_analyzer.py
index 055de96a5e5..39decb75206 100644
--- a/envs/monkey_zoo/blackbox/analyzers/communication_analyzer.py
+++ b/envs/monkey_zoo/blackbox/analyzers/communication_analyzer.py
@@ -1,4 +1,5 @@
-from typing import Iterable
+from ipaddress import IPv4Address
+from typing import Collection, Iterable
 
 from envs.monkey_zoo.blackbox.analyzers.analyzer import Analyzer
 from envs.monkey_zoo.blackbox.analyzers.analyzer_log import AnalyzerLog
@@ -13,15 +14,22 @@ def __init__(self, island_client: MonkeyIslandClient, machine_ips: Iterable[str]
 
     def analyze_test_results(self):
         self.log.clear()
-        all_monkeys_communicated = True
+        all_agents_communicated = True
+        agent_ips = self._get_agent_ips()
+
         for machine_ip in self.machine_ips:
-            if not self.did_monkey_communicate_back(machine_ip):
-                self.log.add_entry("Monkey from {} didn't communicate back".format(machine_ip))
-                all_monkeys_communicated = False
+            if self._agent_communicated_back(machine_ip, agent_ips):
+                self.log.add_entry("Agent from {} communicated back".format(machine_ip))
             else:
-                self.log.add_entry("Monkey from {} communicated back".format(machine_ip))
-        return all_monkeys_communicated
+                self.log.add_entry("Agent from {} didn't communicate back".format(machine_ip))
+                all_agents_communicated = False
+
+        return all_agents_communicated
+
+    def _get_agent_ips(self) -> Collection[IPv4Address]:
+        agents = self.island_client.get_agents()
+        machines = self.island_client.get_machines()
+        return {i.ip for a in agents for i in machines[a.machine_id].network_interfaces}
 
-    def did_monkey_communicate_back(self, machine_ip: str):
-        query = {"ip_addresses": {"$elemMatch": {"$eq": machine_ip}}}
-        return len(self.island_client.find_monkeys_in_db(query)) > 0
+    def _agent_communicated_back(self, machine_ip: str, agent_ips: Collection[IPv4Address]) -> bool:
+        return IPv4Address(machine_ip) in agent_ips