"""
Faraday Penetration Test IDE
Copyright (C) 2013  Infobyte LLC (http://www.infobytesec.com/)
See the file 'doc/LICENSE' for the license information
"""
import re
from collections import defaultdict

from faraday_plugins.plugins.plugin import PluginBase

__author__ = "Francisco Amato"
__copyright__ = "Copyright (c) 2013, Infobyte LLC"
__credits__ = ["Francisco Amato"]
__license__ = ""
__version__ = "1.0.0"
__maintainer__ = "Francisco Amato"
__email__ = "famato@infobytesec.com"
__status__ = "Development"


class DnsmapParser:
    """
    The objective of this class is to parse an xml file generated by the
    dnsmap tool.

    TODO: Handle errors.
    TODO: Test dnsmap output version. Handle what happens if the parser
    doesn't support it.
    TODO: Test cases.

    @param dnsmap_filepath A proper simple report generated by dnsmap
    """

    def __init__(self, output):
        self.items = defaultdict(list)
        if "\n\n" in output:
            self.parse_txt(output)
        else:
            self.parse_csv(output)

    def parse_txt(self, output):
        hosts = self.split_output_lines(output)

        for host_data in hosts:
            if len(host_data) == 2:
                ip = self.clean_ip(host_data[1])
                hostname = host_data[0]
                self.add_host_info_to_items(ip, hostname)
            elif len(host_data) > 2:
                hostname = host_data.pop(0)
                for ip_address in host_data:
                    ip = self.clean_ip(ip_address)
                    self.add_host_info_to_items(ip, hostname)

    def parse_csv(self, output):
        hosts = list(filter(None, output.splitlines()))

        for host in hosts:
            host_data = host.split(",", 1)
            if host_data[1].count(',') == 0:
                ip = host_data[1]
                hostname = host_data[0]
                self.add_host_info_to_items(ip, hostname)
            else:
                hostname = host_data.pop(0)
                ips = host_data[0].split(",")
                for ip_address in ips:
                    self.add_host_info_to_items(ip_address, hostname)

    def split_output_lines(self, output):
        splitted = output.splitlines()
        hosts_list = []
        aux_list = []
        for i in range(0, len(splitted)):
            if not splitted[i]:
                hosts_list.append(aux_list)
                aux_list = []
                continue
            else:
                aux_list.append(splitted[i])
        return hosts_list

    def clean_ip(self, item):
        ip = item.split(':', 1)
        return ip[1].strip()

    def add_host_info_to_items(self, ip_address, hostname):
        self.items[ip_address].append(hostname)


class DnsmapPlugin(PluginBase):
    """Example plugin to parse dnsmap output."""

    def __init__(self, *arg, **kwargs):
        super().__init__(*arg, **kwargs)
        self.id = "Dnsmap"
        self.name = "Dnsmap Output Plugin"
        self.plugin_version = "0.3"
        self.version = "0.30"
        self.options = None
        self._current_output = None
        self._command_regex = re.compile(r'^(sudo dnsmap|dnsmap|\.\/dnsmap)\s+.*?')
        self.xml_arg_re = re.compile(r"^.*(-r\s*[^\s]+).*$")
        self._use_temp_file = True
        self._temp_file_extension = "txt"

    def canParseCommandString(self, current_input):
        if self._command_regex.match(current_input.strip()):
            self.command = self.get_command(current_input)
            return True
        else:
            return False

    def parseOutputString(self, output):
        """
        This method will discard the output the shell sends, it will read it
        from the xml where it expects it to be present.
        """
        parser = DnsmapParser(output)
        for ip_address, hostnames in parser.items.items():
            h_id = self.createAndAddHost(ip_address, hostnames=hostnames)
        return True

    def processCommandString(self, username, current_path, command_string):
        """
        Adds the parameter to get output to the command string that the
        user has set.
        """
        super().processCommandString(username, current_path, command_string)
        arg_match = self.xml_arg_re.match(command_string)

        if arg_match is None:
            return f"{command_string} -r {self._output_file_path} \\n"
        else:
            return re.sub(arg_match.group(1),
                          r"-r %s" % self._output_file_path,
                          command_string)


def createPlugin(*args, **kwargs):
    return DnsmapPlugin(*args, **kwargs)
