Refactor state telem

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:

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}