PBA's stored on the database

monkey/infection_monkey/monkey.py

@@ -113,13 +113,7 @@ class InfectionMonkey(object):
             system_info = system_info_collector.get_info()
             ControlClient.send_telemetry("system_info_collection", system_info)
 
-        pb = PostBreach()
+        PostBreach().execute()
-        output = pb.execute()
-        ControlClient.send_telemetry("post_breach", {'output': output})
-
-        for action_class in WormConfiguration.post_breach_actions:
-            action = action_class()
-            action.act()
 
         if 0 == WormConfiguration.depth:
             LOG.debug("Reached max depth, shutting down")

monkey/infection_monkey/post_breach/post_breach.py

@@ -2,6 +2,7 @@ import logging
 import infection_monkey.config
 import subprocess
 import platform
+from infection_monkey.control import ControlClient
 
 LOG = logging.getLogger(__name__)
 
@@ -15,27 +16,41 @@ class PostBreach(object):
 
     def execute(self):
         for pba in self.pba_list:
-            if platform.system() == 'Windows':
+            pba.run()
-                return pba.execute_win()
-            else:
-                return pba.execute_linux()
 
     @staticmethod
     def config_to_pba_list(config):
         """
-        Should return a list of PBA's generated from config
+        Should return a list of PBA's generated from config. After ATT&CK is implemented this will pick
+        which PBA's to run.
         """
         pba_list = []
-        if config.post_breach_actions["linux"] or config.post_breach_actions["windows"]:
+        # Get custom PBA command from config
-            pba_list.append(PBA(config.post_breach_actions["linux"], config.post_breach_actions["windows"]))
+        custom_pba_linux = config.post_breach_actions['linux'] if "linux" in config.post_breach_actions else ""
+        custom_pba_windows = config.post_breach_actions['windows'] if "windows" in config.post_breach_actions else ""
+
+        if custom_pba_linux or custom_pba_windows:
+            pba_list.append(PBA('custom_pba', custom_pba_linux, custom_pba_windows))
         return pba_list
 
 
 # Post Breach Action container
 class PBA(object):
-    def __init__(self, linux_command="", windows_command=""):
+    def __init__(self, name="unknown", linux_command="", windows_command=""):
         self.linux_command = linux_command
         self.windows_command = windows_command
+        self.name = name
+
+    def run(self):
+        if platform.system() == 'Windows':
+            ControlClient.send_telemetry('post_breach', {'command': self.windows_command,
+                                                         'output': self.execute_win(),
+                                                         'name': self.name})
+        else:
+            ControlClient.send_telemetry('post_breach', {'command': self.linux_command,
+                                                         'output': self.execute_linux(),
+                                                         'name': self.name})
+        return False
 
     def execute_linux(self):
         return subprocess.check_output(self.linux_command, shell=True) if self.linux_command else False

monkey/monkey_island/cc/resources/telemetry.py

@@ -259,11 +259,9 @@ class Telemetry(flask_restful.Resource):
 
     @staticmethod
     def process_post_breach_telemetry(telemetry_json):
-        if telemetry_json['output']:
+        mongo.db.monkey.update(
-            node = NodeService.get_or_create_node(telemetry_json['ip'], telemetry_json['domain_name'])
+            {'guid': telemetry_json['monkey_guid']},
-
+            {'$push': {'post_breach_actions': telemetry_json['data']}})
-
-        pass
 
 TELEM_PROCESS_DICT = \
     {

monkey/monkey_island/cc/resources/telemetry_feed.py

@@ -80,6 +80,15 @@ class TelemetryFeed(flask_restful.Resource):
     def get_trace_telem_brief(telem):
         return 'Monkey reached max depth.'
 
+    @staticmethod
+    def get_post_breach_telem_brief(telem):
+        target = telem['data']['ip']
+        output = telem['data']['output']
+        if output:
+            return 'Monkey ran post breach commands on %s.' % target
+        else:
+            return 'Monkey failed running post breach commands on %s.' % target
+
 
 TELEM_PROCESS_DICT = \
     {
@@ -88,5 +97,6 @@ TELEM_PROCESS_DICT = \
         'exploit': TelemetryFeed.get_exploit_telem_brief,
         'scan': TelemetryFeed.get_scan_telem_brief,
         'system_info_collection': TelemetryFeed.get_systeminfo_telem_brief,
-        'trace': TelemetryFeed.get_trace_telem_brief
+        'trace': TelemetryFeed.get_trace_telem_brief,
+        'post_breach': TelemetryFeed.get_post_breach_telem_brief
     }