Merge pull request #48 from mitre/emu-overwrite

Avoid emu data overwrite

Author: deacon-mp

app/emu_svc.py

@@ -159,8 +159,11 @@ def _copy_planner(self, source_path, target_filename):
         if not os.path.exists(planner_dir):
             os.makedirs(planner_dir)
         target_path = os.path.join(planner_dir, target_filename)
-        shutil.copyfile(source_path, target_path)
-        self.log.debug('Copied planner to %s', target_path)
+        if os.path.exists(target_path):
+            self.log.debug('Planner %s already exists. Skipping.', target_path)
+        else:
+            shutil.copyfile(source_path, target_path)
+            self.log.debug('Copied planner to %s', target_path)
 
     @staticmethod
     def _is_planner(data):
@@ -226,8 +229,11 @@ async def _write_adversary(self, data):
             os.makedirs(d)
 
         file_path = os.path.join(d, '%s.yml' % data['id'])
-        with open(file_path, 'w') as f:
-            f.write(yaml.dump(data))
+        if os.path.exists(file_path):
+            self.log.debug('Adversary profile %s already exists. Skipping.', file_path)
+        else:
+            with open(file_path, 'w') as f:
+                f.write(yaml.dump(data))
 
     async def _save_adversary(self, id, name, description, abilities):
         adversary = dict(
@@ -249,8 +255,11 @@ async def _write_ability(self, data):
         if not os.path.exists(d):
             os.makedirs(d)
         file_path = os.path.join(d, '%s.yml' % data['id'])
-        with open(file_path, 'w') as f:
-            f.write(yaml.dump([data]))
+        if os.path.exists(file_path):
+            self.log.debug('Ability file %s already exists. Skipping.', file_path)
+        else:
+            with open(file_path, 'w') as f:
+                f.write(yaml.dump([data]))
 
     @staticmethod
     def get_privilege(executors):
@@ -344,5 +353,8 @@ async def _write_source(self, data):
             os.makedirs(d)
 
         file_path = os.path.join(d, '%s.yml' % data['id'])
-        with open(file_path, 'w') as f:
-            f.write(yaml.dump(data))
+        if os.path.exists(file_path):
+            self.log.debug('Fact source file %s already exists. Skipping.', file_path)
+        else:
+            with open(file_path, 'w') as f:
+                f.write(yaml.dump(data))

hook.py

@@ -1,5 +1,4 @@
 import os
-import shutil
 
 from app.utility.base_world import BaseWorld
 from plugins.emu.app.emu_svc import EmuService
@@ -22,10 +21,5 @@ async def enable(services):
     if not os.path.isdir(plugin_svc.repo_dir):
         await plugin_svc.clone_repo()
 
-    for directory in ['abilities', 'adversaries', 'sources', 'planners']:
-        full_path = os.path.join(data_dir, directory)
-        if os.path.isdir(full_path):
-            shutil.rmtree(full_path)
-
     await plugin_svc.decrypt_payloads()
     await plugin_svc.populate_data_directory()