import os
import re
from io import BytesIO

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

"""
from lxml import etree
from lxml.etree import XMLParser
from faraday_plugins.plugins.plugin import PluginXMLFormat

current_path = os.path.abspath(os.getcwd())


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

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

    @param nmap_xml_filepath A proper xml generated by nmap
    """

    def __init__(self, xml_output):
        tree = self.parse_xml(xml_output)

        if tree:
            self.hosts = [host for host in self.get_hosts(tree)]
        else:
            self.hosts = []

    def parse_xml(self, xml_output):
        """
        Open and parse an xml file.

        TODO: Write custom parser to just read the nodes that we need instead
         of reading the whole file.

        @return xml_tree An xml tree instance. None if error.
        """

        try:
            magical_parser = XMLParser(recover=True)
            return etree.parse(BytesIO(xml_output), magical_parser)
        except SyntaxError as err:
            # logger.error("SyntaxError: %s." % (err))
            return None

    def get_hosts(self, tree):
        """
        @return hosts A list of Host instances
        """
        for host_node in tree.findall('host'):
            yield Host(host_node)


def get_attrib_from_subnode(xml_node, subnode_xpath_expr, attrib_name):
    """
    Finds a subnode in the host node and the retrieves a value from it

    @return An attribute value
    """

    node = xml_node.find(subnode_xpath_expr)

    if node is not None:
        return node.get(attrib_name)

    return None


class Host:
    """
    An abstract representation of a Host

    TODO: Consider evaluating the attributes lazily
    TODO: Write what's expected to be present in the nodes
    TODO: Refactor both Host and the Port clases?

    @param host_node A host_node taken from an nmap xml tree
    """

    def __init__(self, host_node):
        self.node = host_node

        self.hostnames = [hostname[0] for hostname in self.get_hostnames()]
        if len(self.hostnames) != 0:
            self.hostname = self.hostnames[0]
        else:
            self.hostname = 'unknown'

        self.hostnames = list(set(self.hostnames))
        self.status = self.get_status()
        self.ipv4_address = self.get_ipv4_address()
        self.ipv6_address = self.get_ipv6_address()
        self.mac_address = self.get_mac_address()
        self.os_guesses = [os_guess for os_guess in self.get_os_guesses()]
        self.os = self.top_os_guess()
        self.ports = [port for port in self.get_ports()]
        self.vulns = [vuln for vuln in self.get_scripts()]
        if self.os != 'unknown':
            for p in self.ports:
                if p.service is not None:
                    if p.service.ostype:
                        self.os = p.service.ostype
                        break

    def get_hostnames(self):
        """
        Expects to find one or more
        '<hostname name="localhost.localdomain" type="PTR"/>' in the host node.

        @return A list of (hostname, hostname_type) or None
        """
        for hostname in self.node.findall('hostnames/hostname'):
            yield (hostname.attrib["name"], hostname.attrib["type"])

    def get_attrib_from_subnode(self, subnode_xpath_expr, attrib_name):
        """
        Finds a subnode in the host node and the retrieves a value from it

        @return An attribute value
        """
        return get_attrib_from_subnode(
            self.node,
            subnode_xpath_expr,
            attrib_name)

    def get_status(self):
        """
        Expects to find '<status state="up" reason="conn-refused"/>'
        in the node
        TODO: Use 'reason'
        @return An status or 'unknown'
        """
        status = self.get_attrib_from_subnode('status', 'state')

        return status if status else 'unknown'

    def get_ipv4_address(self):
        """
        Expects to find '<address addr="127.0.0.1" addrtype="ipv4"/>'
        in the node

        @return ip_address or 'unknown'
        """
        ip_address = self.get_attrib_from_subnode(
            "address[@addrtype='ipv4']",
            'addr')
        return ip_address if ip_address else 'unknown'

    def get_ipv6_address(self):
        """
        Expects to find '<address addr="127.0.0.1" addrtype="ipv6"/>'
        in the node

        @return ip_address or 'unknown'
        """
        ip_address = self.get_attrib_from_subnode(
            "address[@addrtype='ipv6']",
            'addr')
        return ip_address if ip_address else 'unknown'

    def get_mac_address(self):
        """
        Expects to find
        '<address addr="00:08:54:26:A9:E5" addrtype="mac" vendor="Netronix" />'
        in the node

        @return mac_address or 'unknown'
        """
        mac_address = self.get_attrib_from_subnode(
            "address[@addrtype='mac']",
            'addr')
        return mac_address if mac_address else 'unknown'

    def get_os_guesses(self):
        """
        Expects to find
        '<os>..<osclass type="general purpose" vendor="Microsoft"
        osfamily="Windows" osgen="2003" accuracy="96" />..</os>' in the node

        @return A list of (os_vendor_family_gen, accuracy)
        """
        # OS information about host with great acurracy.

        osclasses = self.node.findall('os/osclass')
        if osclasses == []:
            osclasses = self.node.findall('os/osmatch/osclass')

        for osclass in osclasses:
            os_vendor = osclass.get("vendor", "unknown")
            os_family = osclass.get("osfamily", "unknown")
            os_gen = osclass.get("osgen", "unknown")
            accuracy = osclass.get("accuracy", "unknown")

            yield (f"{os_vendor} {os_family} {os_gen}", accuracy)

        # Os information in services, bad acurracy.
        if osclasses == []:
            services = self.node.findall("ports/port/service")
            for service in services:
                ostype = service.get("ostype", "unknown")
                yield (f"{ostype}", 0)

    def top_os_guess(self):
        """
        @return The most accurate os_guess_id or 'unknown'.
        """
        return self.os_guesses[0][0] if len(self.os_guesses) != 0 else 'unknown'

    def get_scripts(self):
        # Expects to find a scripts in the node.
        for s in self.node.findall('hostscript/script'):
            yield Script(s)

    def get_ports(self):
        """
        Expects to find one or more
        '<port protocol="tcp" portid="631">...</port>' in the node.

        @return A list of Port instances or None
        """
        for port in self.node.findall('ports/port'):
            yield Port(port)

    def is_up(self):
        """
        Returns True if the host is up else False.
        """
        if self.status == 'up':
            return True
        else:
            return False

    def __str__(self):
        ports = []
        for port in self.ports:
            var = f"    {port}"
            ports.append(var)
        ports = "\n".join(ports)

        return f"{self.hostnames}, {self.status}, {self.ipv4_address} [{self.mac_address}], {self.os}\n{ports}"


class Port:
    """
    An abstract representation of a Port.

    @param port_node A port_node taken from an nmap xml tree
    """

    PORT_STATUS_FIX = {"filtered": "closed", "open|filtered": "closed"}

    def __init__(self, port_node):
        self.node = port_node

        self.protocol = self.node.get("protocol")
        self.number = self.node.get("portid")
        self.state, self.reason, self.reason_ttl = self.get_state()
        self.service = self.get_service()
        self.vulns = [vuln for vuln in self.get_scripts()]

    def get_attrib_from_subnode(self, subnode_xpath_expr, attrib_name):
        """
        Finds a subnode in the host node and the retrieves a value from it.

        @return An attribute value
        """
        return get_attrib_from_subnode(
            self.node,
            subnode_xpath_expr,
            attrib_name)

    def get_state(self):
        """
        Expects to find a
        '<state state="open" reason="syn-ack" reason_ttl="0"/>' in the node.

        @return (state, reason, reason_ttl) or ('unknown','unknown','unknown')
        """
        state = self.PORT_STATUS_FIX.get(self.get_attrib_from_subnode('state', 'state'),
                                         self.get_attrib_from_subnode('state', 'state'))
        reason = self.get_attrib_from_subnode('state', 'reason')
        reason_ttl = self.get_attrib_from_subnode('state', 'reason_ttl')

        return (state if state else 'unknown',
                reason if reason else 'unknown',
                reason_ttl if reason_ttl else 'unknown')

    def get_service(self):
        """
        Expects to find a service in the node.
        """
        service_node = self.node.find('service')
        if service_node is not None:
            return Service(service_node)

        return None

    def get_scripts(self):
        """
        Expects to find a scripts in the node.
        """
        for s in self.node.findall('script'):
            if s.get("id") == 'vulners':
                for service_table in s.findall('table'):
                    for vulnerability_table in service_table.findall('table'):
                        yield ScriptVulners(s, service_table, vulnerability_table)
            else:
                yield Script(s)

    def __str__(self):
        return f"{self.number}, {self.state}, Service: {self.service}"


class ScriptVulners:
    """
    An abstract representation of a script table when script is 'vulners'.
    https://nmap.org/nsedoc/scripts/vulners.html

    '<script id="vulners" output="&#xa;  cpe:/a:apache:http_server:2.4.6: ...">
    <table key="cpe:/a:apache:http_server:2.4.6">
    <table><elem key="cvss">7.5</elem>
    <elem key="id">CVE-2017-7679</elem><elem key="type">cve</elem>
    <elem key="is_exploit">false</elem></table>
    ...
    </table></script>'

    @param script_node The reference to the node of the 'vulners' script
    @param service_table The outer table taken from an nmap 'vulners' script xml tree
    @param vulnerability_table The inner table taken from an nmap 'vulners' script xml tree
    """

    def __init__(self, script_node, service_table, vulnerability_table):
        self.node = script_node
        self.table = {}
        for e in vulnerability_table.findall('elem'):
            self.table[e.get("key")] = str(e.text)

        self.name = self.table["id"]

        self.desc = script_node.get("id") + "-" + self.table["id"]
        if self.table["is_exploit"] == 'true':
            self.desc += " *EXPLOIT*"

        self.refs = ["https://vulners.com/" + self.table["type"] + "/" + self.table["id"]]
        self.response = ""
        self.web = ""
        self.cvss2 = {}

    def __str__(self):
        return f"{self.name}, {self.product}, {self.version}"


class Script:
    """
    An abstract representation of a Script.

    '<script id="http-methods" output="No Allow or Public header in OPTIONS
    response (status code 400)"/><script id="http-title"
    output="Document Error: Unauthorized"><elem key="title">
    Document Error: Unauthorized</elem></script>'

    @param script_node A script_node taken from an nmap xml tree
    """

    def parse_output(self, output):
        block_re = re.compile('^\\s{4}References:((?:.|[\r\n])+[\r\n](?:\\s{4}\\w|\\s*$))', re.MULTILINE)
        m1 = block_re.findall(output)
        if len(m1) > 0:
            links_re = re.compile('[ \t]+([^ \t\n\r]+)[ \t]*')
            m2 = links_re.findall(m1[0])
            return m2
        return []

    def __init__(self, script_node):
        self.node = script_node

        self.name = script_node.get("id")
        self.desc = script_node.get("output")
        self.refs = self.parse_output(self.desc)
        self.response = ""
        for k in script_node.findall("elem"):
            self.response += "\n" + str(k.get('key')) + ": " + str(k.text)
        self.web = re.search("(http-|https-)", self.name)
        self.cvss2 = {}

    def __str__(self):
        return f"{self.name}, {self.product}, {self.version}"


class Service:
    """
    An abstract representation of a Service.

    '<service name="ipp" product="CUPS" version="1.4" method="probed"
     conf="10"/>'

    @param service_node A service_node taken from an nmap xml tree
    """

    def __init__(self, service_node):
        self.node = service_node

        name = service_node.get("name")
        self.name = name if name else 'unknown'

        self.tunnel = service_node.get("tunnel")
        if self.tunnel == "ssl":
            if self.name == "http":
                self.name = "https"
            elif self.name == 'imap':
                self.name = 'imaps'
            elif self.name == 'pop3':
                self.name = 'pop3s'

        product = service_node.get("product")
        self.product = product if product else 'unknown'

        version = service_node.get("version")
        self.version = version if version else 'unknown'

        self.method = service_node.get("method")
        self.conf = service_node.get("conf")
        self.ostype = self.node.get("ostype")

    def __str__(self):
        return f"{self.name}, {self.product}, {self.version}"


class NmapPlugin(PluginXMLFormat):
    """
    Example plugin to parse nmap output.
    """

    def __init__(self, *arg, **kwargs):
        super().__init__(*arg, **kwargs)
        self.identifier_tag = "nmaprun"
        self.id = "Nmap"
        self.name = "Nmap XML Output Plugin"
        self.plugin_version = "0.0.3"
        self.version = "6.40"
        self.framework_version = "1.0.0"
        self.options = None
        self._current_output = None
        self._command_regex = re.compile(r'^(sudo nmap|nmap|\.\/nmap|sudo masscan|masscan)\s+.*?')
        self._use_temp_file = True
        self._temp_file_extension = "xml"
        self.xml_arg_re = re.compile(r"^.*(-oX\s*[^\s]+).*$")
        self.addSetting("Scan Technique", str, "-sS")

    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.

        NOTE: if 'debug' is true then it is being run from a test case and the
        output being sent is valid.
        """

        parser = NmapXmlParser(output)

        for host in parser.hosts:
            # if not host.is_up():
            #     continue

            if host.mac_address == 'unknown':
                host.mac_address = "00:00:00:00:00:00"

            if host.ipv4_address != 'unknown':
                minterfase = host.ipv4_address
            else:
                minterfase = host.ipv6_address
            h_id = self.createAndAddHost(minterfase, host.os, mac=host.mac_address, hostnames=host.hostnames)

            for v in host.vulns:
                desc = v.desc
                desc += "\nOutput: " + v.response if v.response else ""

                v_id = self.createAndAddVulnToHost(
                    h_id,
                    v.name,
                    desc=v.desc,
                    ref=v.refs,
                    severity=0,
                    cve=[v.name]
                )

            for port in host.ports:

                srvname = str(port.number)
                srvversion = "unknown"
                if port.service is not None:
                    srvname = port.service.name
                    srvversion = port.service.product if port.service.product != "unknown" else ""
                    srvversion += " " + port.service.version if port.service.version != "unknown" else ""

                s_id = self.createAndAddServiceToHost(
                    h_id,
                    srvname,
                    port.protocol,
                    ports=[port.number],
                    status=port.state,
                    version=srvversion,
                    description=srvname)

                for v in port.vulns:

                    desc = v.desc
                    refs = v.refs

                    if hasattr(v, 'severity'):
                        severity = v.severity
                    else:
                        severity = "info"
                        if re.search(r"(?<!NOT )VULNERABLE", desc):
                            severity = "high"
                        if re.search(r"ERROR", desc):
                            severity = "unclassified"
                        if re.search(r"Couldn't", desc):
                            severity = "unclassified"

                    if v.web:
                        v_id = self.createAndAddVulnWebToService(
                            h_id,
                            s_id,
                            v.name,
                            desc=desc,
                            response=v.response if v.response else "",
                            ref=refs,
                            severity=severity,
                            website=minterfase,
                            cve=[v.name],
                            cvss2=v.cvss2
                        )
                    else:
                        v_id = self.createAndAddVulnToService(
                            h_id,
                            s_id,
                            v.name,
                            desc=v.desc,
                            ref=refs,
                            severity=severity,
                            cve=[v.name],
                            cvss2=v.cvss2
                        )
        del parser

    def processCommandString(self, username, current_path, command_string):
        """
        Adds the -oX parameter to get xml 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:
            if "masscan" in command_string:
                return re.sub(r"(^.*?masscan)",
                              r"\1 -oX %s" % self._output_file_path,
                              command_string)
            else:
                return re.sub(r"(^.*?nmap)",
                              r"\1 -oX %s" % self._output_file_path,
                              command_string)
        else:
            return re.sub(arg_match.group(1),
                          r"-oX %s" % self._output_file_path,
                          command_string)


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