From 27ca921dbcb42cac4251a2bb119094e91122bec4 Mon Sep 17 00:00:00 2001
From: itay <itay.mizeretz@guardicore.com>
Date: Sun, 23 Jun 2019 15:36:28 +0300
Subject: [PATCH] Refactor state telem

---
 monkey/infection_monkey/monkey.py             |  5 +++--
 .../infection_monkey/telemetry/state_telem.py | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)
 create mode 100644 monkey/infection_monkey/telemetry/state_telem.py

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}