diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 06e0f267b..9c1ffb419 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -16,6 +16,7 @@ from infection_monkey.network.network_scanner import NetworkScanner from infection_monkey.system_info import SystemInfoCollector from infection_monkey.system_singleton import SystemSingleton from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem +from infection_monkey.telemetry.state_telem import StateTelem from infection_monkey.windows_upgrader import WindowsUpgrader from infection_monkey.post_breach.post_breach_handler import PostBreach from common.utils.attack_utils import ScanStatus @@ -109,7 +110,7 @@ class InfectionMonkey(object): if monkey_tunnel: monkey_tunnel.start() - ControlClient.send_telemetry("state", {'done': False}) + StateTelem(False).send() self._default_server = WormConfiguration.current_server LOG.debug("default server: %s" % self._default_server) @@ -223,7 +224,7 @@ class InfectionMonkey(object): InfectionMonkey.close_tunnel() firewall.close() else: - ControlClient.send_telemetry("state", {'done': True}) # Signal the server (before closing the tunnel) + StateTelem(False).send() # Signal the server (before closing the tunnel) InfectionMonkey.close_tunnel() firewall.close() if WormConfiguration.send_log_to_server: diff --git a/monkey/infection_monkey/telemetry/state_telem.py b/monkey/infection_monkey/telemetry/state_telem.py new file mode 100644 index 000000000..3bd63d2f9 --- /dev/null +++ b/monkey/infection_monkey/telemetry/state_telem.py @@ -0,0 +1,19 @@ +from infection_monkey.telemetry.base_telem import BaseTelem + +__author__ = "itay.mizeretz" + + +class StateTelem(BaseTelem): + + def __init__(self, is_done): + """ + Default state telemetry constructor + :param is_done: Whether the state of monkey is done. + """ + super(StateTelem, self).__init__() + self.is_done = is_done + + telem_category = 'state' + + def get_data(self): + return {'done': self.is_done}