Merge pull request #364 from guardicore/feature/refactor-telem-retro

Feature/refactor telem retro

monkey/infection_monkey/exploit/__init__.py

@@ -20,33 +20,30 @@ class HostExploiter(object):
 
     def __init__(self, host):
         self._config = infection_monkey.config.WormConfiguration
-        self._exploit_info = {'display_name': self._EXPLOITED_SERVICE,
+        self.exploit_info = {'display_name': self._EXPLOITED_SERVICE,
                               'started': '',
                               'finished': '',
                               'vulnerable_urls': [],
                               'vulnerable_ports': [],
                               'executed_cmds': []}
-        self._exploit_attempts = []
+        self.exploit_attempts = []
         self.host = host
 
     def set_start_time(self):
-        self._exploit_info['started'] = datetime.now().isoformat()
+        self.exploit_info['started'] = datetime.now().isoformat()
 
     def set_finish_time(self):
-        self._exploit_info['finished'] = datetime.now().isoformat()
+        self.exploit_info['finished'] = datetime.now().isoformat()
 
     def is_os_supported(self):
         return self.host.os.get('type') in self._TARGET_OS_TYPE
 
     def send_exploit_telemetry(self, result):
-        from infection_monkey.control import ControlClient
-        ControlClient.send_telemetry(
-            'exploit',
-            {'result': result, 'machine': self.host.__dict__, 'exploiter': self.__class__.__name__,
-             'info': self._exploit_info, 'attempts': self._exploit_attempts})
+        from infection_monkey.telemetry.exploit_telem import ExploitTelem
+        ExploitTelem(self, result).send()
 
     def report_login_attempt(self, result, user, password='', lm_hash='', ntlm_hash='', ssh_key=''):
-        self._exploit_attempts.append({'result': result, 'user': user, 'password': password,
+        self.exploit_attempts.append({'result': result, 'user': user, 'password': password,
                                        'lm_hash': lm_hash, 'ntlm_hash': ntlm_hash, 'ssh_key': ssh_key})
 
     def exploit_host(self):
@@ -66,10 +63,10 @@ class HostExploiter(object):
         raise NotImplementedError()
 
     def add_vuln_url(self, url):
-        self._exploit_info['vulnerable_urls'].append(url)
+        self.exploit_info['vulnerable_urls'].append(url)
 
     def add_vuln_port(self, port):
-        self._exploit_info['vulnerable_ports'].append(port)
+        self.exploit_info['vulnerable_ports'].append(port)
 
     def add_executed_cmd(self, cmd):
         """
@@ -77,7 +74,7 @@ class HostExploiter(object):
         :param cmd: String of executed command. e.g. 'echo Example'
         """
         powershell = True if "powershell" in cmd.lower() else False
-        self._exploit_info['executed_cmds'].append({'cmd': cmd, 'powershell': powershell})
+        self.exploit_info['executed_cmds'].append({'cmd': cmd, 'powershell': powershell})
 
 
 from infection_monkey.exploit.win_ms08_067 import Ms08_067_Exploiter

monkey/infection_monkey/exploit/sambacry.py

@@ -65,9 +65,9 @@ class SambaCryExploiter(HostExploiter):
         LOG.info("Writable shares and their credentials on host %s: %s" %
                  (self.host.ip_addr, str(writable_shares_creds_dict)))
 
-        self._exploit_info["shares"] = {}
+        self.exploit_info["shares"] = {}
         for share in writable_shares_creds_dict:
-            self._exploit_info["shares"][share] = {"creds": writable_shares_creds_dict[share]}
+            self.exploit_info["shares"][share] = {"creds": writable_shares_creds_dict[share]}
             self.try_exploit_share(share, writable_shares_creds_dict[share])
 
         # Wait for samba server to load .so, execute code and create result file.
@@ -90,7 +90,7 @@ class SambaCryExploiter(HostExploiter):
             self.clean_share(self.host.ip_addr, share, writable_shares_creds_dict[share])
 
         for share, fullpath in successfully_triggered_shares:
-            self._exploit_info["shares"][share]["fullpath"] = fullpath
+            self.exploit_info["shares"][share]["fullpath"] = fullpath
 
         if len(successfully_triggered_shares) > 0:
             LOG.info(

monkey/infection_monkey/exploit/shellshock.py

@@ -66,7 +66,7 @@ class ShellShockExploiter(HostExploiter):
         exploitable_urls = [url for url in exploitable_urls if url[0] is True]
 
         # we want to report all vulnerable URLs even if we didn't succeed
-        self._exploit_info['vulnerable_urls'] = [url[1] for url in exploitable_urls]
+        self.exploit_info['vulnerable_urls'] = [url[1] for url in exploitable_urls]
 
         # now try URLs until we install something on victim
         for _, url, header, exploit in exploitable_urls:

monkey/infection_monkey/exploit/weblogic.py

@@ -108,7 +108,7 @@ class WebLogic201710271(WebRCE):
         if httpd.get_requests > 0:
             # Add all urls because we don't know which one is vulnerable
             self.vulnerable_urls.extend(urls)
-            self._exploit_info['vulnerable_urls'] = self.vulnerable_urls
+            self.exploit_info['vulnerable_urls'] = self.vulnerable_urls
         else:
             LOG.info("No vulnerable urls found, skipping.")
 

monkey/infection_monkey/monkey.py

@@ -16,6 +16,11 @@ 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.system_info_telem import SystemInfoTelem
+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
@@ -109,24 +114,23 @@ 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)
-        ControlClient.send_telemetry("tunnel", {'proxy': ControlClient.proxies.get('https')})
+        TunnelTelem().send()
 
         if WormConfiguration.collect_system_info:
             LOG.debug("Calling system info collection")
             system_info_collector = SystemInfoCollector()
             system_info = system_info_collector.get_info()
-            ControlClient.send_telemetry("system_info_collection", system_info)
+            SystemInfoTelem(system_info).send()
 
         # Executes post breach actions
         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)
@@ -157,8 +161,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:
@@ -223,7 +226,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/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,17 +46,7 @@ class PBA(object):
         """
         exec_funct = self._execute_default
         result = exec_funct()
-        try:
-            hostname = socket.gethostname()
-            ip = socket.gethostbyname(hostname)
-        except socket.error:
-            hostname = "Unknown"
-            ip = "Unknown"
-        ControlClient.send_telemetry('post_breach', {'command': self.command,
-                                                     'result': result,
-                                                     'name': self.name,
-                                                     'hostname': hostname,
-                                                     'ip': ip})
+        PostBreachTelem(self, result).send()
 
     def _execute_default(self):
         """

monkey/infection_monkey/telemetry/exploit_telem.py

@@ -0,0 +1,27 @@
+from infection_monkey.telemetry.base_telem import BaseTelem
+
+__author__ = "itay.mizeretz"
+
+
+class ExploitTelem(BaseTelem):
+
+    def __init__(self, exploiter, result):
+        """
+        Default exploit telemetry constructor
+        :param exploiter: The instance of exploiter used
+        :param result: The result from the 'exploit_host' method.
+        """
+        super(ExploitTelem, self).__init__()
+        self.exploiter = exploiter
+        self.result = result
+
+    telem_category = 'exploit'
+
+    def get_data(self):
+        return {
+            'result': self.result,
+            'machine': self.exploiter.host.__dict__,
+            'exploiter': self.exploiter.__class__.__name__,
+            'info': self.exploiter.exploit_info,
+            'attempts': self.exploiter.exploit_attempts
+        }

monkey/infection_monkey/telemetry/post_breach_telem.py

@@ -0,0 +1,40 @@
+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, self.ip = PostBreachTelem._get_hostname_and_ip()
+
+    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
+        }
+
+    @staticmethod
+    def _get_hostname_and_ip():
+        try:
+            hostname = socket.gethostname()
+            ip = socket.gethostbyname(hostname)
+        except socket.error:
+            hostname = "Unknown"
+            ip = "Unknown"
+        return hostname, 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/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}

monkey/infection_monkey/telemetry/system_info_telem.py

@@ -0,0 +1,19 @@
+from infection_monkey.telemetry.base_telem import BaseTelem
+
+__author__ = "itay.mizeretz"
+
+
+class SystemInfoTelem(BaseTelem):
+
+    def __init__(self, system_info):
+        """
+        Default system info telemetry constructor
+        :param system_info: System info returned from SystemInfoCollector.get_info()
+        """
+        super(SystemInfoTelem, self).__init__()
+        self.system_info = system_info
+
+    telem_category = 'system_info'
+
+    def get_data(self):
+        return self.system_info

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.py

@@ -79,7 +79,7 @@ class Telemetry(flask_restful.Resource):
                 monkey_label = telem_monkey_guid
             x["monkey"] = monkey_label
             objects.append(x)
-            if x['telem_category'] == 'system_info_collection' and 'credentials' in x['data']:
+            if x['telem_category'] == 'system_info' and 'credentials' in x['data']:
                 for user in x['data']['credentials']:
                     if -1 != user.find(','):
                         new_user = user.replace(',', '.')
@@ -277,7 +277,7 @@ TELEM_PROCESS_DICT = \
         'state': Telemetry.process_state_telemetry,
         'exploit': Telemetry.process_exploit_telemetry,
         'scan': Telemetry.process_scan_telemetry,
-        'system_info_collection': Telemetry.process_system_info_telemetry,
+        'system_info': Telemetry.process_system_info_telemetry,
         'trace': Telemetry.process_trace_telemetry,
         'post_breach': Telemetry.process_post_breach_telemetry,
         'attack': Telemetry.process_attack_telemetry

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):
@@ -97,7 +97,7 @@ TELEM_PROCESS_DICT = \
         'state': TelemetryFeed.get_state_telem_brief,
         'exploit': TelemetryFeed.get_exploit_telem_brief,
         'scan': TelemetryFeed.get_scan_telem_brief,
-        'system_info_collection': TelemetryFeed.get_systeminfo_telem_brief,
+        'system_info': TelemetryFeed.get_systeminfo_telem_brief,
         'trace': TelemetryFeed.get_trace_telem_brief,
         'post_breach': TelemetryFeed.get_post_breach_telem_brief,
         'attack': TelemetryFeed.get_attack_telem_brief

monkey/monkey_island/cc/services/report.py

@@ -172,7 +172,7 @@ class ReportService:
         PASS_TYPE_DICT = {'password': 'Clear Password', 'lm_hash': 'LM hash', 'ntlm_hash': 'NTLM hash'}
         creds = []
         for telem in mongo.db.telemetry.find(
-                {'telem_category': 'system_info_collection', 'data.credentials': {'$exists': True}},
+                {'telem_category': 'system_info', 'data.credentials': {'$exists': True}},
                 {'data.credentials': 1, 'monkey_guid': 1}
         ):
             monkey_creds = telem['data']['credentials']
@@ -200,7 +200,7 @@ class ReportService:
         """
         creds = []
         for telem in mongo.db.telemetry.find(
-                {'telem_category': 'system_info_collection', 'data.ssh_info': {'$exists': True}},
+                {'telem_category': 'system_info', 'data.ssh_info': {'$exists': True}},
                 {'data.ssh_info': 1, 'monkey_guid': 1}
         ):
             origin = NodeService.get_monkey_by_guid(telem['monkey_guid'])['hostname']
@@ -221,7 +221,7 @@ class ReportService:
         """
         creds = []
         for telem in mongo.db.telemetry.find(
-                {'telem_category': 'system_info_collection', 'data.Azure': {'$exists': True}},
+                {'telem_category': 'system_info', 'data.Azure': {'$exists': True}},
                 {'data.Azure': 1, 'monkey_guid': 1}
         ):
             azure_users = telem['data']['Azure']['usernames']
@@ -383,7 +383,7 @@ class ReportService:
     @staticmethod
     def get_monkey_subnets(monkey_guid):
         network_info = mongo.db.telemetry.find_one(
-            {'telem_category': 'system_info_collection', 'monkey_guid': monkey_guid},
+            {'telem_category': 'system_info', 'monkey_guid': monkey_guid},
             {'data.network_info.networks': 1}
         )
         if network_info is None: