diff --git a/envs/monkey_zoo/blackbox/analyzers/communication_analyzer.py b/envs/monkey_zoo/blackbox/analyzers/communication_analyzer.py index 055de96a5..39decb752 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 @@ class CommunicationAnalyzer(Analyzer): def analyze_test_results(self): self.log.clear() - all_monkeys_communicated = True - 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 - else: - self.log.add_entry("Monkey from {} communicated back".format(machine_ip)) - return all_monkeys_communicated + all_agents_communicated = True + agent_ips = self._get_agent_ips() - 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 + for machine_ip in self.machine_ips: + 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("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 _agent_communicated_back(self, machine_ip: str, agent_ips: Collection[IPv4Address]) -> bool: + return IPv4Address(machine_ip) in agent_ips