Bugfix in dropper.py, handle gracefully failure in cleanup

gitlabsyncint · Apr 17, 2018 · 558fa74 · 558fa74
1 parent cc4ad05
commit 558fa74
Showing 1 changed file with 21 additions and 18 deletions.
diff --git a/infection_monkey/dropper.py b/infection_monkey/dropper.py
@@ -133,22 +133,25 @@ def start(self):
             LOG.warn("Seems like monkey died too soon")
 
     def cleanup(self):
-        if (self._config['source_path'].lower() != self._config['destination_path'].lower()) and \
-                os.path.exists(self._config['source_path']) and \
-                WormConfiguration.dropper_try_move_first:
+        try:
+            if (self._config['source_path'].lower() != self._config['destination_path'].lower()) and \
+                    os.path.exists(self._config['source_path']) and \
+                    WormConfiguration.dropper_try_move_first:
 
-            # try removing the file first
-            try:
-                os.remove(self._config['source_path'])
-            except Exception as exc:
-                LOG.debug("Error removing source file '%s': %s", self._config['source_path'], exc)
-
-                # mark the file for removal on next boot
-                dropper_source_path_ctypes = c_char_p(self._config['source_path'])
-                if 0 == ctypes.windll.kernel32.MoveFileExA(dropper_source_path_ctypes, None,
-                                                           MOVEFILE_DELAY_UNTIL_REBOOT):
-                    LOG.debug("Error marking source file '%s' for deletion on next boot (error %d)",
-                              self._config['source_path'], ctypes.windll.kernel32.GetLastError())
-                else:
-                    LOG.debug("Dropper source file '%s' is marked for deletion on next boot",
-                              self._config['source_path'])
+                # try removing the file first
+                try:
+                    os.remove(self._config['source_path'])
+                except Exception as exc:
+                    LOG.debug("Error removing source file '%s': %s", self._config['source_path'], exc)
+
+                    # mark the file for removal on next boot
+                    dropper_source_path_ctypes = c_char_p(self._config['source_path'])
+                    if 0 == ctypes.windll.kernel32.MoveFileExA(dropper_source_path_ctypes, None,
+                                                               MOVEFILE_DELAY_UNTIL_REBOOT):
+                        LOG.debug("Error marking source file '%s' for deletion on next boot (error %d)",
+                                  self._config['source_path'], ctypes.windll.kernel32.GetLastError())
+                    else:
+                        LOG.debug("Dropper source file '%s' is marked for deletion on next boot",
+                                  self._config['source_path'])
+        except AttributeError:
+            LOG.error("Invalid configuration options. Failing")