Implemented file restoration endpoint, file upload front end

monkey/monkey_island/cc/app.py

@@ -120,7 +120,8 @@ def init_app(mongo_url):
     api.add_resource(IslandLog, '/api/log/island/download', '/api/log/island/download/')
     api.add_resource(PBAFileDownload, '/api/pba/download/<string:path>')
     api.add_resource(FileUpload, '/api/fileUpload/<string:file_type>',
-                     '/api/fileUpload/<string:file_type>?load=<string:file_name>')
+                     '/api/fileUpload/<string:file_type>?load=<string:filename>',
+                     '/api/fileUpload/<string:file_type>?restore=<string:filename>')
     api.add_resource(RemoteRun, '/api/remote-monkey', '/api/remote-monkey/')
 
     return app

monkey/monkey_island/cc/resources/file_upload.py

@@ -4,25 +4,35 @@ from cc.services.config import ConfigService
 import os
 from werkzeug.utils import secure_filename
 import logging
+from cc.database import mongo
+import copy
 
 LOG = logging.getLogger(__name__)
 UPLOADS_DIR = "./monkey_island/cc/userUploads"
+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):
-    # Used by monkey. can't secure.
-    def get(self, path):
-        return send_from_directory(UPLOADS_DIR, path)
 
-    def get(self, file_type, file_name):
+    def get(self, file_type):
-        req_data = request.data
+        # Verify that file_name is indeed a file from config
+        if file_type == LINUX_PBA_TYPE:
+            filename = ConfigService.get_config_value(copy.deepcopy(LINUX_PBA_INFO))['name']
+        else:
+            filename = ConfigService.get_config_value(copy.deepcopy(WINDOWS_PBA_INFO))['name']
+        return send_from_directory(GET_FILE_DIR, filename)
 
     def post(self, file_type):
-        filename = ''
+        filename = FileUpload.upload_pba_file(request, (file_type == LINUX_PBA_TYPE))
-        if file_type == 'PBAlinux':
-            filename = FileUpload.upload_pba_file(request)
-        elif file_type == 'PBAwindows':
-            filename = FileUpload.upload_pba_file(request, False)
 
         response = Response(
             response=filename,
@@ -30,32 +40,40 @@ class FileUpload(flask_restful.Resource):
         return response
 
     def delete(self, file_type):
-        config = ConfigService.get_config(should_decrypt=False)
+        file_conf_path = LINUX_PBA_INFO if file_type == 'PBAlinux' else WINDOWS_PBA_INFO
-        if file_type == 'PBAlinux':
+        filename = ConfigService.get_config_value(file_conf_path)['name']
-            file_info = 'linux_file_info'
-        else:
-            file_info = 'windows_file_info'
-        filename = config['monkey']['behaviour']['custom_post_breach'][file_info]['name']
         file_path = os.path.join(UPLOADS_DIR, filename)
         try:
             if os.path.exists(file_path):
                 os.remove(file_path)
+            ConfigService.set_config_value(file_conf_path, {'size': '0', 'name': ''})
         except OSError as e:
             LOG.error("Can't remove previously uploaded post breach files: %s" % e)
+
         return {}
 
     @staticmethod
     def upload_pba_file(request_, is_linux=True):
-        config = ConfigService.get_config(should_decrypt=False)
         filename = secure_filename(request_.files['filepond'].filename)
         file_path = os.path.join(UPLOADS_DIR, filename)
         request_.files['filepond'].save(file_path)
         file_size = os.path.getsize(file_path)
+        ConfigService.set_config_value((LINUX_PBA_INFO if is_linux else WINDOWS_PBA_INFO),
+                                       {'size': file_size, 'name': filename})
+        return filename
+
+    @staticmethod
+    def get_file_info_db_paths(is_linux=True):
+        """
+        Gets PBA file size and name parameter config paths for linux and windows
+        :param is_linux: determines whether to get windows or linux file params
+        :return: returns tuple of filename and file size paths in config
+        """
         if is_linux:
             file_info = 'linux_file_info'
         else:
             file_info = 'windows_file_info'
-        config['monkey']['behaviour']['custom_post_breach'][file_info]['size'] = file_size
+        config_path = 'monkey.behaviour.custom_post_breach.' + file_info + '.'
-        config['monkey']['behaviour']['custom_post_breach'][file_info]['name'] = filename
+        size_path = config_path + 'size'
-        ConfigService.update_config(config, should_encrypt=False)
+        name_path = config_path + 'name'
-        return filename
+        return name_path, size_path

monkey/monkey_island/cc/services/config.py

@@ -85,6 +85,12 @@ class ConfigService:
                 config = encryptor.dec(config)
         return config
 
+    @staticmethod
+    def set_config_value(config_key_as_arr, value):
+        mongo_key = ".".join(config_key_as_arr)
+        mongo.db.config.update({'name': 'newconfig'},
+                               {"$set": {mongo_key: value}})
+
     @staticmethod
     def get_flat_config(is_initial_config=False, should_decrypt=True):
         config_json = ConfigService.get_config(is_initial_config, should_decrypt)

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

@@ -78,6 +78,7 @@ class AppComponent extends AuthComponent {
   constructor(props) {
     super(props);
     this.state = {
+      removePBAfiles: false,
       completedSteps: {
         run_server: true,
         run_monkey: false,
@@ -88,6 +89,11 @@ class AppComponent extends AuthComponent {
     };
   }
 
+  // Sets the property that indicates if we need to remove PBA files from state or not
+  setRemovePBAfiles = (rmFiles) => {
+    this.setState({removePBAfiles: rmFiles});
+  };
+
   componentDidMount() {
     this.updateStatus();
     this.interval = setInterval(this.updateStatus, 5000);
@@ -173,11 +179,15 @@ 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

@@ -3,7 +3,7 @@ import Form from 'react-jsonschema-form';
 import {Col, Nav, NavItem} from 'react-bootstrap';
 import fileDownload from 'js-file-download';
 import AuthComponent from '../AuthComponent';
-import { FilePond, registerPlugin } from 'react-filepond';
+import { FilePond } from 'react-filepond';
 import 'filepond/dist/filepond.min.css';
 
 class ConfigurePageComponent extends AuthComponent {
@@ -21,7 +21,9 @@ class ConfigurePageComponent extends AuthComponent {
       lastAction: 'none',
       sections: [],
       selectedSection: 'basic',
-      allMonkeysAreDead: true
+      allMonkeysAreDead: true,
+      PBAwinFile: [],
+      PBAlinuxFile: []
     };
   }
 
@@ -95,6 +97,7 @@ class ConfigurePageComponent extends AuthComponent {
   };
 
   resetConfig = () => {
+    this.removePBAfiles();
     this.authFetch('/api/configuration/island',
       {
         method: 'POST',
@@ -112,6 +115,16 @@ 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')){
+      this.PBAlinuxPond.removeFile();
+      this.PBAwindowsPond.removeFile();
+    }
+    this.setState({PBAlinuxFile: []});
+    this.setState({PBAwinFile: []});
+  }
+
   onReadFile = (event) => {
     try {
       this.setState({
@@ -153,16 +166,75 @@ class ConfigurePageComponent extends AuthComponent {
   };
 
   PBAwindows = () => {
-    if (! this.state.configuration.monkey.behaviour.custom_post_breach.windows_file_info){
+    return (<FilePond
-      
+      server='/api/fileUpload/PBAwindows'
-    }
+      files={this.getWinPBAfile()}
-    return (<FilePond server='/api/fileUpload/PBAwindows'/>)
+      onupdatefiles={fileItems => {
+        this.setState({
+          PBAwinFile: fileItems.map(fileItem => fileItem.file)
+        })
+      }}
+      ref={ref => this.PBAwindowsPond = ref}
+    />)
   };
 
   PBAlinux = () => {
-    return (<FilePond server='/api/fileUpload/PBAlinux'/>)
+    return (<FilePond
+      server='/api/fileUpload/PBAlinux'
+      files={this.getLinuxPBAfile()}
+      onupdatefiles={fileItems => {
+        this.setState({
+          PBAlinuxFile: fileItems.map(fileItem => fileItem.file)
+        })
+      }}
+      ref={ref => this.PBAlinuxPond = ref}
+      onload={this.props.setRemovePBAfiles(false)}
+    />)
+
   };
 
+  getWinPBAfile(){
+    if (this.props.removePBAfiles){
+      // 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 {
+      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 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 {
+      console.log("Getting from config")
+      return ConfigurePageComponent.getPBAfile(this.state.configuration.monkey.behaviour.custom_post_breach.linux_file_info)
+    }
+  }
+
+  static getPBAfile(fileSrc, isMock=false){
+    let PBAfile = [{
+      source: fileSrc.name,
+      options: {
+        type: 'limbo'
+      }
+    }];
+    if (isMock){
+      PBAfile[0].options.file = fileSrc
+    }
+    return PBAfile
+  }
 
   render() {
     let displayedSchema = {};
@@ -180,6 +252,12 @@ class ConfigurePageComponent extends AuthComponent {
           },
           windows_file: {
             "ui:widget": this.PBAwindows
+          },
+          linux_file_info: {
+            classNames: "linux-pba-file-info"
+          },
+          windows_file_info: {
+            classNames: "windows-pba-file-info"
           }
         }
       }
@@ -188,7 +266,6 @@ class ConfigurePageComponent extends AuthComponent {
       displayedSchema = this.state.schema['properties'][this.state.selectedSection];
       displayedSchema['definitions'] = this.state.schema['definitions'];
     }
-
     return (
       <Col xs={12} lg={8}>
         <h1 className="page-title">Monkey Configuration</h1>
@@ -276,7 +353,6 @@ class ConfigurePageComponent extends AuthComponent {
             </div>
             : ''}
         </div>
-
       </Col>
     );
   }

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

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

monkey/monkey_island/cc/ui/src/styles/App.css

@@ -163,6 +163,18 @@ body {
  * Configuration Page
  */
 
+.linux-pba-file-info, .windows-pba-file-info {
+  display: none
+}
+
+.filepond--root li {
+  overflow: visible;
+}
+
+.filepond--root * {
+  font-size: 1.04em;
+}
+
 .rjsf .form-group .form-group {
   margin-left: 2em;
 }