PR notes fixed, command + file bugs fixed

monkey/infection_monkey/post_breach/file_execution.py

@@ -1,8 +1,8 @@
 from infection_monkey.post_breach.pba import PBA
 from infection_monkey.control import ControlClient
 from infection_monkey.config import WormConfiguration
+from infection_monkey.utils import get_monkey_dir_path
 import requests
-import shutil
 import os
 import logging
 
@@ -25,27 +25,25 @@ class FileExecution(PBA):
         super(FileExecution, self).__init__("File execution", linux_command, windows_command)
 
     def _execute_linux(self):
-        FileExecution.download_PBA_file(FileExecution.get_dest_dir(WormConfiguration, True),
+        FileExecution.download_PBA_file(get_monkey_dir_path(), self.linux_filename)
-                                        self.linux_filename)
         return super(FileExecution, self)._execute_linux()
 
     def _execute_win(self):
-        FileExecution.download_PBA_file(FileExecution.get_dest_dir(WormConfiguration, True),
+        FileExecution.download_PBA_file(get_monkey_dir_path(), self.windows_filename)
-                                        self.windows_filename)
         return super(FileExecution, self)._execute_win()
 
     def add_default_command(self, is_linux):
         """
-        Replaces current (likely empty) command with default file execution command.
+        Replaces current (likely empty) command with default file execution command (that changes permissions, executes
+        and finally deletes post breach file).
+        Default commands are defined as globals in this module.
         :param is_linux: Boolean that indicates for which OS the command is being set.
         """
         if is_linux:
-            file_path = os.path.join(FileExecution.get_dest_dir(WormConfiguration, is_linux=True),
+            file_path = os.path.join(get_monkey_dir_path(), self.linux_filename)
-                                     self.linux_filename)
             self.linux_command = DEFAULT_LINUX_COMMAND.format(file_path)
         else:
-            file_path = os.path.join(FileExecution.get_dest_dir(WormConfiguration, is_linux=False),
+            file_path = os.path.join(get_monkey_dir_path(), self.windows_filename)
-                                     self.windows_filename)
             self.windows_command = DEFAULT_WINDOWS_COMMAND.format(file_path)
 
     @staticmethod
@@ -63,18 +61,8 @@ class FileExecution(PBA):
                                          proxies=ControlClient.proxies)
         try:
             with open(os.path.join(dst_dir, filename), 'wb') as written_PBA_file:
-                shutil.copyfileobj(PBA_file_contents, written_PBA_file)
+                written_PBA_file.write(PBA_file_contents.content)
             return True
         except IOError as e:
             LOG.error("Can not download post breach file to target machine, because %s" % e)
             return False
-
-    @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/post_breach_handler.py

@@ -3,11 +3,14 @@ import infection_monkey.config
 from file_execution import FileExecution
 from pba import PBA
 from infection_monkey.utils import is_windows_os
+from infection_monkey.utils import get_monkey_dir_path
 
 LOG = logging.getLogger(__name__)
 
 __author__ = 'VakarisZ'
 
+DIR_CHANGE_WINDOWS = 'cd %s & '
+DIR_CHANGE_LINUX = 'cd %s ; '
 
 class PostBreach(object):
     """
@@ -31,7 +34,6 @@ class PostBreach(object):
         Returns a list of PBA objects generated from config.
         :param config: Monkey configuration
         :return: A list of PBA objects.
-        TODO: Parse PBA's from PBA array (like 'add_user'). Also merge the whole outdated PBA structure into this one.
         """
         pba_list = []
         pba_list.extend(PostBreach.get_custom_PBA(config))
@@ -49,19 +51,23 @@ class PostBreach(object):
         file_pba = FileExecution()
         command_pba = PBA(name="Custom")
 
+        if not is_windows_os():
             # Add linux commands to PBA's
             if config.PBA_linux_filename:
                 if config.custom_PBA_linux_cmd:
-                file_pba.linux_command = config.custom_PBA_linux_cmd
+                    # Add change dir command, because user will try to access his file
+                    file_pba.linux_command = (DIR_CHANGE_LINUX % get_monkey_dir_path()) + config.custom_PBA_linux_cmd
                 else:
                     file_pba.add_default_command(is_linux=True)
             elif config.custom_PBA_linux_cmd:
                 command_pba.linux_command = config.custom_PBA_linux_cmd
-
+        else:
             # Add windows commands to PBA's
             if config.PBA_windows_filename:
                 if config.custom_PBA_windows_cmd:
-                file_pba.windows_command = config.custom_PBA_windows_cmd
+                    # Add change dir command, because user will try to access his file
+                    file_pba.windows_command = (DIR_CHANGE_WINDOWS % get_monkey_dir_path()) + \
+                                               config.custom_PBA_windows_cmd
                 else:
                     file_pba.add_default_command(is_linux=False)
             elif config.custom_PBA_windows_cmd:

monkey/monkey_island/cc/resources/pba_file_download.py

@@ -1,6 +1,6 @@
 import flask_restful
 from flask import send_from_directory
-from cc.services.config import UPLOADS_DIR
+from cc.resources.pba_file_upload import GET_FILE_DIR
 
 __author__ = 'VakarisZ'
 
@@ -11,4 +11,4 @@ class PBAFileDownload(flask_restful.Resource):
     """
     # Used by monkey. can't secure.
     def get(self, path):
-        return send_from_directory(UPLOADS_DIR, path)
+        return send_from_directory(GET_FILE_DIR, path)

monkey/monkey_island/cc/resources/pba_file_upload.py

@@ -1,6 +1,7 @@
 import flask_restful
 from flask import request, send_from_directory, Response
-from cc.services.config import ConfigService, PBA_WINDOWS_FILENAME_PATH, PBA_LINUX_FILENAME_PATH, UPLOADS_DIR
+from cc.services.config import ConfigService
+from cc.services.post_breach_files import PBA_WINDOWS_FILENAME_PATH, PBA_LINUX_FILENAME_PATH, UPLOADS_DIR
 from cc.auth import jwt_required
 import os
 from werkzeug.utils import secure_filename

monkey/monkey_island/cc/resources/root.py

@@ -10,6 +10,7 @@ from cc.services.config import ConfigService
 from cc.services.node import NodeService
 from cc.services.report import ReportService
 from cc.utils import local_ip_addresses
+from cc.services.post_breach_files import remove_PBA_files
 
 __author__ = 'Barak'
 
@@ -42,7 +43,7 @@ class Root(flask_restful.Resource):
     @staticmethod
     @jwt_required()
     def reset_db():
-        ConfigService.remove_PBA_files()
+        remove_PBA_files()
         # We can't drop system collections.
         [mongo.db[x].drop() for x in mongo.db.collection_names() if not x.startswith('system.')]
         ConfigService.init_config()

monkey/monkey_island/cc/services/config.py

@@ -4,7 +4,7 @@ import functools
 import logging
 from jsonschema import Draft4Validator, validators
 from six import string_types
-import os
+import cc.services.post_breach_files
 
 from cc.database import mongo
 from cc.encryptor import encryptor
@@ -34,12 +34,6 @@ ENCRYPTED_CONFIG_STRINGS = \
         ['cnc', 'aws_config', 'aws_secret_access_key']
     ]
 
-UPLOADS_DIR = 'monkey_island/cc/userUploads'
-
-# Where to find file names in config
-PBA_WINDOWS_FILENAME_PATH = ['monkey', 'behaviour', 'PBA_windows_filename']
-PBA_LINUX_FILENAME_PATH = ['monkey', 'behaviour', 'PBA_linux_filename']
-
 
 class ConfigService:
     default_config = None
@@ -152,7 +146,7 @@ class ConfigService:
     @staticmethod
     def update_config(config_json, should_encrypt):
         # PBA file upload happens on pba_file_upload endpoint and corresponding config options are set there
-        ConfigService.keep_PBA_files(config_json)
+        cc.services.post_breach_files.set_config_PBA_files(config_json)
         if should_encrypt:
             try:
                 ConfigService.encrypt_config(config_json)
@@ -163,18 +157,6 @@ class ConfigService:
         logger.info('monkey config was updated')
         return True
 
-    @staticmethod
-    def keep_PBA_files(config_json):
-        """
-        Sets PBA file info in config_json to current config's PBA file info values.
-        :param config_json: config_json that will be modified
-        """
-        if ConfigService.get_config():
-            linux_filename = ConfigService.get_config_value(PBA_LINUX_FILENAME_PATH)
-            windows_filename = ConfigService.get_config_value(PBA_WINDOWS_FILENAME_PATH)
-            config_json['monkey']['behaviour']['PBA_linux_filename'] = linux_filename
-            config_json['monkey']['behaviour']['PBA_windows_filename'] = windows_filename
-
     @staticmethod
     def init_default_config():
         if ConfigService.default_config is None:
@@ -200,7 +182,7 @@ class ConfigService:
 
     @staticmethod
     def reset_config():
-        ConfigService.remove_PBA_files()
+        cc.services.post_breach_files.remove_PBA_files()
         config = ConfigService.get_default_config(True)
         ConfigService.set_server_ips_in_config(config)
         ConfigService.update_config(config, should_encrypt=False)
@@ -309,22 +291,3 @@ class ConfigService:
             pair['public_key'] = encryptor.dec(pair['public_key'])
             pair['private_key'] = encryptor.dec(pair['private_key'])
         return pair
-
-    @staticmethod
-    def remove_PBA_files():
-        if ConfigService.get_config():
-            linux_filename = ConfigService.get_config_value(PBA_WINDOWS_FILENAME_PATH)
-            windows_filename = ConfigService.get_config_value(PBA_LINUX_FILENAME_PATH)
-            if linux_filename:
-                ConfigService.remove_file(linux_filename)
-            if windows_filename:
-                ConfigService.remove_file(windows_filename)
-
-    @staticmethod
-    def remove_file(file_name):
-        file_path = os.path.join(UPLOADS_DIR, file_name)
-        try:
-            if os.path.exists(file_path):
-                os.remove(file_path)
-        except OSError as e:
-            logger.error("Can't remove previously uploaded post breach files: %s" % e)

monkey/monkey_island/cc/services/post_breach_files.py

@@ -0,0 +1,43 @@
+import cc.services.config
+import logging
+import os
+
+__author__ = "VakarisZ"
+
+logger = logging.getLogger(__name__)
+
+# Where to find file names in config
+PBA_WINDOWS_FILENAME_PATH = ['monkey', 'behaviour', 'PBA_windows_filename']
+PBA_LINUX_FILENAME_PATH = ['monkey', 'behaviour', 'PBA_linux_filename']
+UPLOADS_DIR = 'monkey_island/cc/userUploads'
+
+
+def remove_PBA_files():
+    if cc.services.config.ConfigService.get_config():
+        windows_filename = cc.services.config.ConfigService.get_config_value(PBA_WINDOWS_FILENAME_PATH)
+        linux_filename = cc.services.config.ConfigService.get_config_value(PBA_LINUX_FILENAME_PATH)
+        if linux_filename:
+            remove_file(linux_filename)
+        if windows_filename:
+            remove_file(windows_filename)
+
+
+def remove_file(file_name):
+    file_path = os.path.join(UPLOADS_DIR, file_name)
+    try:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+    except OSError as e:
+        logger.error("Can't remove previously uploaded post breach files: %s" % e)
+
+
+def set_config_PBA_files(config_json):
+    """
+    Sets PBA file info in config_json to current config's PBA file info values.
+    :param config_json: config_json that will be modified
+    """
+    if cc.services.config.ConfigService.get_config():
+        linux_filename = cc.services.config.ConfigService.get_config_value(PBA_LINUX_FILENAME_PATH)
+        windows_filename = cc.services.config.ConfigService.get_config_value(PBA_WINDOWS_FILENAME_PATH)
+        config_json['monkey']['behaviour']['PBA_linux_filename'] = linux_filename
+        config_json['monkey']['behaviour']['PBA_windows_filename'] = windows_filename

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

@@ -58,7 +58,7 @@ class PostBreachComponent extends React.Component {
 
   render() {
     let pbaMachines = this.props.data.filter(function(value, index, arr){
-        return ( value.pba_results !== "None" && value.pba_results.length);
+        return ( value.pba_results !== "None" && value.pba_results.length > 0);
     });
     let defaultPageSize = pbaMachines.length > pageSize ? pageSize : pbaMachines.length;
     let showPagination = pbaMachines > pageSize;