Refactor scan,trace,tunnel,pba telems

monkey/infection_monkey/monkey.py

@@ -16,7 +16,10 @@ 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.scan_telem import ScanTelem
 from infection_monkey.telemetry.state_telem import StateTelem
+from infection_monkey.telemetry.trace_telem import TraceTelem
+from infection_monkey.telemetry.tunnel_telem import TunnelTelem
 from infection_monkey.windows_upgrader import WindowsUpgrader
 from infection_monkey.post_breach.post_breach_handler import PostBreach
 from common.utils.attack_utils import ScanStatus
@@ -114,7 +117,7 @@ class InfectionMonkey(object):
 
         self._default_server = WormConfiguration.current_server
         LOG.debug("default server: %s" % self._default_server)
-        ControlClient.send_telemetry("tunnel", {'proxy': ControlClient.proxies.get('https')})
+        TunnelTelem().send()
 
         if WormConfiguration.collect_system_info:
             LOG.debug("Calling system info collection")
@@ -126,8 +129,7 @@ class InfectionMonkey(object):
         PostBreach().execute()
 
         if 0 == WormConfiguration.depth:
-            LOG.debug("Reached max depth, shutting down")
-            ControlClient.send_telemetry("trace", "Reached max depth, shutting down")
+            TraceTelem("Reached max depth, shutting down").send()
             return
         else:
             LOG.debug("Running with depth: %d" % WormConfiguration.depth)
@@ -158,8 +160,7 @@ class InfectionMonkey(object):
                              machine, finger.__class__.__name__)
                     finger.get_host_fingerprint(machine)
 
-                ControlClient.send_telemetry('scan', {'machine': machine.as_dict(),
-                                                      'service_count': len(machine.services)})
+                ScanTelem(machine).send()
 
                 # skip machines that we've already exploited
                 if machine in self._exploited_machines:

monkey/infection_monkey/post_breach/pba.py

@@ -2,6 +2,7 @@ import logging
 import subprocess
 import socket
 from infection_monkey.control import ControlClient
+from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 from infection_monkey.utils import is_windows_os
 from infection_monkey.config import WormConfiguration
 
@@ -45,12 +46,7 @@ class PBA(object):
         """
         exec_funct = self._execute_default
         result = exec_funct()
-        hostname = socket.gethostname()
-        ControlClient.send_telemetry('post_breach', {'command': self.command,
-                                                     'result': result,
-                                                     'name': self.name,
-                                                     'hostname': hostname,
-                                                     'ip': socket.gethostbyname(hostname)})
+        PostBreachTelem(self, result).send()
 
     def _execute_default(self):
         """

monkey/infection_monkey/telemetry/post_breach_telem.py

@@ -0,0 +1,31 @@
+import socket
+
+from infection_monkey.telemetry.base_telem import BaseTelem
+
+__author__ = "itay.mizeretz"
+
+
+class PostBreachTelem(BaseTelem):
+
+    def __init__(self, pba, result):
+        """
+        Default post breach telemetry constructor
+        :param pba: Post breach action which was used
+        :param result: Result of PBA
+        """
+        super(PostBreachTelem, self).__init__()
+        self.pba = pba
+        self.result = result
+        self.hostname = socket.gethostname()
+        self.ip = socket.gethostbyname(self.hostname)
+
+    telem_category = 'post_breach'
+
+    def get_data(self):
+        return {
+            'command': self.pba.command,
+            'result': self.result,
+            'name': self.pba.name,
+            'hostname': self.hostname,
+            'ip': self.ip
+        }

monkey/infection_monkey/telemetry/scan_telem.py

@@ -0,0 +1,22 @@
+from infection_monkey.telemetry.base_telem import BaseTelem
+
+__author__ = "itay.mizeretz"
+
+
+class ScanTelem(BaseTelem):
+
+    def __init__(self, machine):
+        """
+        Default scan telemetry constructor
+        :param machine: Scanned machine
+        """
+        super(ScanTelem, self).__init__()
+        self.machine = machine
+
+    telem_category = 'scan'
+
+    def get_data(self):
+        return {
+            'machine': self.machine.as_dict(),
+            'service_count': len(self.machine.services)
+        }

monkey/infection_monkey/telemetry/trace_telem.py

@@ -0,0 +1,26 @@
+import logging
+
+from infection_monkey.telemetry.base_telem import BaseTelem
+
+__author__ = "itay.mizeretz"
+
+LOG = logging.getLogger(__name__)
+
+
+class TraceTelem(BaseTelem):
+
+    def __init__(self, msg):
+        """
+        Default trace telemetry constructor
+        :param msg: Trace message
+        """
+        super(TraceTelem, self).__init__()
+        self.msg = msg
+        LOG.debug("Trace: %s" % msg)
+
+    telem_category = 'trace'
+
+    def get_data(self):
+        return {
+            'msg': self.msg
+        }

monkey/infection_monkey/telemetry/tunnel_telem.py

@@ -0,0 +1,19 @@
+from infection_monkey.control import ControlClient
+from infection_monkey.telemetry.base_telem import BaseTelem
+
+__author__ = "itay.mizeretz"
+
+
+class TunnelTelem(BaseTelem):
+
+    def __init__(self):
+        """
+        Default tunnel telemetry constructor
+        """
+        super(TunnelTelem, self).__init__()
+        self.proxy = ControlClient.proxies.get('https')
+
+    telem_category = 'tunnel'
+
+    def get_data(self):
+        return {'proxy': self.proxy}

monkey/monkey_island/cc/resources/telemetry_feed.py

@@ -78,7 +78,7 @@ class TelemetryFeed(flask_restful.Resource):
 
     @staticmethod
     def get_trace_telem_brief(telem):
-        return 'Monkey reached max depth.'
+        return 'Trace: %s' % telem['data']['msg']
 
     @staticmethod
     def get_post_breach_telem_brief(telem):