# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.


import re
import os
import socket
import xml.etree.cElementTree as Et

from minion.plugins.base import ExternalProcessPlugin


class NMAPPlugin(ExternalProcessPlugin):
    PLUGIN_NAME = 'Nmap'
    PLUGIN_VERSION = '1.0'
    PLUGIN_WEIGHT = 'medium' # previously set to 'light', nmap requires a significant amount of network resources

    NMAP_EXECUTABLE = 'nmap'

    # Default ports to scan, if not set in the plan
    DEFAULT_PORTS = {
        'TCP': ['21', '22', '80', '443']
    }

    # Default scan type, if not set in the plan (-sT)
    DEFAULT_SCANTYPE = 'tcp_connect'

    # Default severity for ports and versions that don't match any entry in the plan
    DEFAULT_SEVERITY = 'medium'


    def do_configure(self, configuration=None, enable_logging=False, logging_file='/var/log/minion/nmap.log'):
        """
        Initialize the nmap plugin (aka __init__)

        This sets up some internal variables:
            SEVERITY_ORDER: list of severities from info to critical
            SEVERITY_MAPPING: severity level (ie, medium) to its numerical position (ie, 2) for severity comparisons
            DEFAULT_SEVERITY: a default severity to raise for issues if there is neither a port nor version severity

        Args:
            configuration:  allows the forcible setting of configuration (the 'configuration' section of a plugin's plan
                            when not being run through minion directly; very useful for test cases
            enable_logging: if set to True, allows a plugin writer to create debug logs (ala console.log) with
                            self.logger.debug(msg)
            logging_file:   where to store the debug log generated by setting enable_logging to True
        """

        # This is useful for testing, where a new instance can pass in a configuration object
        if configuration: self.configuration = configuration

        self.SEVERITY_ORDER = ['info', 'low', 'medium', 'high', 'critical']

        # 0 -> info, 1 -> low, etc.
        self.SEVERITY_MAPPING = {severity: self.SEVERITY_ORDER.index(severity) for severity in self.SEVERITY_ORDER}

        # Process the configuration, creating mappings of port -> severity and version -> severity
        self._create_port_to_severity_mappings()
        self._create_version_to_severity_mappings()

        # Variables to hold stdout and stderr
        self.stdout = self.stderr = ''

        # Set the default severities, for both ports and versions
        self.DEFAULT_SEVERITY = self.configuration.get('configuration', {}).get('default_severity', self.DEFAULT_SEVERITY)

        # Enable logging during development
        if enable_logging:
            import logging
            self.logger = logging.getLogger('minion-plugin-nmap')
            self.logger.setLevel(logging.DEBUG)

            fh = logging.FileHandler(logging_file)
            self.logger.addHandler(fh)


    def _create_issue_unauthorized_open_port(self, severity, host, port, product):
        """Generate issues for open ports; they will have one of the three formats for descriptions:

        192.168.0.1: open port (515)
        192.168.0.1: open port (443), running: Apache httpd (authorized software)
        192.168.0.1: open port (445), running: Samba smbd 3.X (unrecognized software)

        Args:
            severity: a dictionary containing both type (port, version, etc.), as well as the severity level (low, etc.)
                      as well as nominal other data (recognized/unrecognized for version scans, etc.)
            host:     hostname/IP address associated with the issue
            port:     open port number associated with issue
            product:  product name (such as OpenSSH 6.2) associated with the issue; may also be None

        Returns:
            A dictionary containing an issue that will get passed up to minion-backend.  It should look like:

            {
                'Severity': 'High',
                'Summary': '10.0.1.1: open port (88), running: Heimdal Kerberos (unrecognized software)',
                'Description': '10.0.1.1: open port (88), running: Heimdal Kerberos (unrecognized software)',
                'URLs': [{'URL': '10.0.1.1:88'}],
                'Ports': [88],
                'Classification': {
                    'cwe_id': '200',
                    'cwe_url': 'http://cwe.mitre.org/data/definitions/200.html'
                }
            }

        Note that cwe_id 200 corresponds to Information Exposure
        """

        sev = severity['severity']
        if sev == 'info':
            authorization = "authorized"
        else:
            authorization = "unauthorized"

        # Generate summaries of the issues
        if 'version' in severity['type']:
            # If software is unrecognized (not matching anything in a severity list), we should say so that instead
            if severity['recognized'] == False:
                authorization = "unrecognized"

            summary = '{host}: open port ({port}), running: {product} ({authorization} software)'.format(
                authorization=authorization, host=host, port=port, product=product)
        elif product:
            summary = '{host}: open port ({port}), running: {product}'.format(
                host=host, port=port, product=product)
        else:
            summary = '{host}: open port ({port})'.format(
                host=host, port=port)

        # Get the severity of the port
        issue = {
            'Severity': sev.capitalize(),
            'Summary': summary,
            'Description': summary,
            'URLs': [{'URL': '{host}:{port}'.format(host=host, port=port)}],
            'Ports': [port],
            'Classification': {
                'cwe_id': '200',
                'cwe_url': 'http://cwe.mitre.org/data/definitions/200.html'
            }
        }

        return issue


    def _create_port_to_severity_mappings(self):
        """
        Generates the internal structure SEVERITY_TO_PORTS, which contains all the levels of severity, along with
        which ports are bound to that severity.  Should be called when the plugin is initialized in do_configure().

        SEVERITY_TO_PORTS should look something like:

        {
            'info': [22],
            'low': [80, 443],
            'medium': [8080],
            'high': [20, 21],
            'critical': [3306]
        }
        """

        self.SEVERITY_TO_PORTS = {}

        for level in self.SEVERITY_ORDER:
            self.SEVERITY_TO_PORTS[level] = set()
            ports = self.configuration.get('severity', {}).get('ports', {}).get(level, [])
            for port in ports:
                self.SEVERITY_TO_PORTS[level].update( self._generate_port_range(port) )


    def _create_version_to_severity_mappings(self):
        """
        Generate the internal structure SEVERITY_TO_VERSION, which contains all the levels of severity, along with
        which regular expressions are bound to that severity.  Note that all regular expressions are both bound
        to the beginning of line (^), and are case-insensitive as well.

        SEVERITY_TO_VERSION should look something like:
        {
            'info': [<_sre.SRE_PATTERN object>],
            'low': [<_sre.SRE_PATTERN object>, <_sre.SRE_PATTERN object>],
            'medium': [<_sre.SRE_PATTERN object>],
            'high': [<_sre.SRE_PATTERN object>, <_sre.SRE_PATTERN object>],
            'critical': [<_sre.SRE_PATTERN object>]
        }
        """

        self.SEVERITY_TO_VERSION = {}

        for level in self.SEVERITY_ORDER:
            self.SEVERITY_TO_VERSION[level] = []
            for version in self.configuration.get('severity', {}).get('version', {}).get(level, []):
                self.SEVERITY_TO_VERSION[level].append(re.compile('^' + version, re.IGNORECASE))


    @staticmethod
    def _generate_port_range(port):
        """
        Generates an array to create a port range from as well as validate that the ports are valid

        Args:
            port: a string containing a port or port range

        Returns:
            an array containing all the port numbers found in that range
            '21' -> [21] or '6666-6668' -> [6666, 6667, 6668]

        Raises:
           Exception: When the ports are invalid -- either non-integers, < 1, or > 65535
        """

        # Cast it to a string first, in case somebody accidentally puts it in their plan as an integer
        port = str(port)

        try:
            # Make sure that every port is an integer, and within the proper range of 1-65535
            port = [int(x) for x in port.split('-')]
            if port[0] < 1 or port[-1] > 65535:
                raise ValueError

            return range(port[0], port[-1] + 1)
        except:
            raise Exception('Invalid port or port range in configuration: "{0}"'.format(port))


    def _generate_switches(self):
        """
        Based on the values set in the configuration -> scan section, generate all the switches necessary to
        run nmap

        Returns:
            An array containing all of the switches necessary to run nmap.  By default, this will include
            -oX, -, --no-stylesheet, -Pn, and --open.
        """

        # We want XML output to stdout, with no stylesheet, and we want to skip host discovery
        switches = ['-oX', '-', '--no-stylesheet', '-Pn', '--open']

        # Mapping of all the friendly names to their nmap scan flags
        scan_type_mappings = {
            'christmas': 'X',
            'fin': 'F',
            'null': 'N',
            'syn': 'S',
            'tcp_connect': 'T',
            'udp': 'U',
            'version': 'V'
        }

        # TODO: check for incompatible flag combinations (low priority)

        # Compile all scan types into a single -s switch, like -sTUV
        scan_type_switch = '-s'
        for scan_type in self.configuration.get('scan', {}).get('types', [self.DEFAULT_SCANTYPE]):
            scan_type_switch += scan_type_mappings.get(scan_type, '')
        switches.append(scan_type_switch)

        # Either we want top ports, or we want to manually specify them
        top_ports = int(self.configuration.get('scan', {}).get('ports', {}).get('top_ports', 0))
        if (top_ports > 0):
            switches += ['--top-ports', str(top_ports)]
        else:
            ports = ''

            # Retrieve the settings from the configuration settings
            tcp = self.configuration.get('scan', {}).get('ports', {}).get('TCP', [])
            udp = self.configuration.get('scan', {}).get('ports', {}).get('UDP', [])

            # Use the default settings if no ports configuration object exists
            if not tcp and not udp:
                tcp = self.DEFAULT_PORTS['TCP']

            # Be lenient in case somebody accidentally` defines a port using an integer instead of a string
            stringify = lambda(x): [str(x) for x in x]  # now I just need to icecreamify
            tcp, udp = stringify(tcp), stringify(udp)

            # Validate every entry
            for entry in tcp + udp:
                self._generate_port_range(entry)

            if (tcp or udp):
                switches.append('-p')
            if tcp:
                ports += 'T:' + ','.join(tcp)

            # Only add UDP ports to the scan if UDP was actually specified in the scan type
            if udp and 'udp' in self.configuration.get('scan', {}).get('types', []):
                if ports: ports += ',U:' + ','.join(udp)
                else: ports = 'U:' + ','.join(udp)

            switches.append(ports)

        return switches


    def get_severity(self, port, product=None):
        """Given a port number and possible product version, return the severity associated with them.

        If both the port and product have a severity associated with it, the plugin will set the severity to the higher
        of the two, unless version_severity_override has been set in the plan.

        If an open port fails to match a value in severity -> port or in severity -> version, then return
        configuration -> default_severity.
        """

        result = {
            'severity': None,
            'type': []
        }

        # Call the severity functions, to get the severities for the passed in ports and products
        port_result = self.get_severity_for_port(port)
        version_result = self.get_severity_for_version(product)

        # Get the severity for both the port and the version
        port_sev = port_result['severity']
        version_sev = version_result['severity']

        # In the case of an item appearing in neither the port nor version list, we just return the internally set
        # default port severity
        if port_sev == None and version_sev == None:
            result['severity'] = self.DEFAULT_SEVERITY
            result['type'] = ['port']

        # There's a defined port severity, but no version severity
        elif port_sev != None and version_sev == None:
            result['severity'] = port_sev
            result['type'] = ['port']

        # There's a defined version severity, but no port severity
        elif port_sev == None and version_sev != None:
            result['severity'] = version_sev
            result['recognized'] = version_result['recognized'] # TODO switch to something more generic?
            result['type'] = ['version']

        # If there's both a port and version severity, we the higher severity of the two, unless we've set
        # the version_severity_override option, in which case we always return the version severity
        else:
            ver_override = self.configuration.get('configuration', {}).get('version_severity_override', False)

            if ver_override:
                result['severity'] = version_sev

            else:
                sev = max(self.SEVERITY_MAPPING[port_sev], self.SEVERITY_MAPPING[version_sev])
                result['severity'] = self.SEVERITY_ORDER[sev]

            result['recognized'] = version_result['recognized']
            result['type'] = ['port', 'version']


        return result


    def get_severity_for_port(self, port):
        """Given a particular port, return the severity of associated with that port. If it fails to match, severity
        will be set to None.

        Args:
            port: integer containing the open port number

        Returns:
            A dictionary containing the severity as well as all whether the product was recognized (ie, it matches
        an entry in severity -> version), or if it wasn't.  For example:

        {'severity': 'high',
         'recognized': False

        """

        result = {
            'severity': None,
        }

        port = int(port)
        for severity in self.SEVERITY_ORDER[::-1]:   # higher priorities take precedence, in case of multiple matches
            if port in self.SEVERITY_TO_PORTS[severity]:
                result['severity'] = severity
                return result

        return result


    def get_severity_for_version(self, product):
        """Given a particular version string (such as OpenSSL 6.2), return the severity of that version.  If the product
        is not recognized, and configuration -> raise_unrecognized_software is set, it will return the severity set in
        configuration -> raise_unrecognized_software_severity.  Otherwise, it will return with a severity of None.

        Args:
            product: The version string generated by nmap, which should consist of product name combined with its version

        Returns:
            A dictionary containing the severity as well as all whether the product was recognized (ie, it matches
        an entry in severity -> version), or if it wasn't.  For example:

        {'severity': 'high',
         'recognized': False}
        """

        result = {
            'severity': None,
            'recognized': False
        }

        if product == None:
            return result

        for severity in self.SEVERITY_ORDER[::-1]:   # higher priorities take precedence, in case of multiple matches
            for expression in self.SEVERITY_TO_VERSION[severity]:
                if expression.search(product) != None:
                    result['recognized'] = True
                    result['severity'] = severity

                    return result

        # If the user has set raise_unrecognized_software, we raise to a specified severity
        if self.configuration.get('configuration', {}).get('raise_unrecognized_software', False):
            result['recognized'] = False
            result['severity'] = self.configuration.get('configuration', {}).get(
                'raise_unrecognized_software_severity', self.DEFAULT_SEVERITY)

        return result


    def generate_issues(self, output):
        """Generates a list of issues based on the parsed nmap output generated by parse_nmap_xml()

        Args:
           output: the dictionary generated by parse_nmap_xml()

        Returns:
           An array of issues, generated from discovered open ports.  See _create_issue_unauthorized_open_port for
           the format of the issue.
        """
        issues = []

        # Sort the list of hosts by IP address; if that fails, then just sort alphabetically
        try:
            hosts = sorted(output.keys(), key=lambda host: socket.inet_aton(host))
        except socket.error:
            hosts = sorted(output.keys())

        # Look up for each host
        for host in hosts:
             for port in sorted(output[host]['ports'].keys(), key=int):
                 entry = output[host]['ports'][port]
                 if 'open' in entry['state']:
                     severity = self.get_severity(port, entry['product'])
                     issues.append(self._create_issue_unauthorized_open_port(severity, host, port, entry['product']))


        return issues


    def parse_nmap_xml(self, xml, baseline={}):
        """Parses the output of nmap -oX --no-stylesheet and returns a Python dictionary containing information on
        each host that was scanned.

        Args:
            xml: A string containing the entirety of the nmap output
            baseline: a JSON object that contains entries that, should they match, are removed from the parsed output

        Returns:
            A dictionary mapping nmap results on each host scanned.  For example:

            '10.0.1.1': {'hostnames': [],
              'ip': '10.0.1.1',
              'ports': {'22': {'name': 'ssh',
                               'product': 'OpenSSH',
                               'protocol': 'tcp',
                               'state': 'open',
                               'version': '6.2'}}},
            '10.0.1.4': {'hostnames': [],
              'ip': '10.0.1.4',
              'ports': {'22': {'name': 'ssh',
                               'product': 'OpenSSH',
                               'protocol': 'tcp',
                               'state': 'open',
                               'version': '6.2'},
                         '443': {'name': 'http',
                                 'product': 'Apache httpd',
                                 'protocol': 'tcp',
                                 'state': 'open',
                                 'version': None},
                         '8080': {'name': 'http',
                                 'product': 'XBMC Web Media Manager',
                                 'protocol': 'tcp',
                                 'state': 'open',
                                 'version': None}}}}

        Raises:
            Exception: If cElementTree is unable to parse the XML output.
        """

        hosts = {}

        try:
            root = Et.fromstring(xml)
        except:
            raise Exception("Unable to parse nmap output")

        # Get every host found
        for host in root.findall("host"):
            # Get the IP
            ip = host.find('address').get('addr')

            # Get the hostname if defined (user selected), otherwise use the IP address as the hostname
            hostname = ''
            hostnames = []
            for hn in host.find('hostnames'):
                if hn.get('type') == 'user':
                    hostname = hn.get('name')
                else:
                    hostnames.append(hn.get('name'))
            if not hostname:
                hostname = ip

            hosts[hostname] = {}
            hosts[hostname]['hostnames'] = hostnames
            hosts[hostname]['ip'] = ip
            hosts[hostname]['ports'] = {}

            # Get open ports, and stuff their data into our hosts data structure
            ports = host.find('ports')
            for port in ports.findall('port'):
                portid = port.get('portid')
                state = port.find('state').get('state')

                service = port.find('service')
                product = service.get('product')
                version = service.get('version')

                if product and version:
                    product = product + ' ' + version

                if 'open' in state:
                    hosts[hostname]['ports'][portid] = {
                        'protocol': port.get('protocol'),
                        'state': state,
                        'product': product,
                        'name': service.get('name'),
                    }


        # If a baseline has been passed in, we need to prune the output of all those entries
        if baseline:
            # Let's remove the __ALLHOSTS__ entry from the baseline, turn into things we can use
            allhosts = baseline.pop('__ALLHOSTS__', {'ports': [], 'products': []})
            allhosts_ports = set()
            for port in allhosts['ports']:
                allhosts_ports.update(self._generate_port_range(port))
            allhosts_products = [re.compile('^' + x, re.IGNORECASE) for x in allhosts['products']]

            # "Deserialize" the baseline, where you can't use periods in key entries (hostnames)
            for baseline_host in baseline.iterkeys():
                baseline[ baseline_host.replace('_', '.') ] = baseline.pop(baseline_host)


            for host in hosts.iterkeys():
                for port in hosts[host]['ports'].keys():
                    # If a port or product is in the __ALLHOSTS__ entry, it shouldn't show up in the results
                    if (int(port) in allhosts_ports or
                        [True for x in allhosts_products if x.search(str(hosts[host]['ports'][port]['product']))]):
                        del(hosts[host]['ports'][port])
                        continue

                    # If a host -> port -> state and host -> port -> product match their mirror in the baseline, remove them
                    # from the output
                    if ((hosts[host]['ports'][port]['state'] ==
                             baseline.get(host, {})
                                     .get('ports', {})
                                     .get(port, {})
                                     .get('state', None)) and
                            (re.search('^' + str(
                            baseline.get(host, {})
                                    .get('ports', {})
                                    .get(port, {})
                                    .get('product', '')), str(hosts[host]['ports'][port]['product']), re.IGNORECASE))):
                        del(hosts[host]['ports'][port])

        return hosts


    def do_start(self):
        """
        Locates and launches the nmap executable against a target that looks either like http://www.website.com,
        192.168.0.0, or 192.168.0.0/24.

        Raises:
            Exception: if the nmap binary is not found somewhere in the system path
        """

        # Locate nmap on the file system
        nmap_path = self.locate_program(self.NMAP_EXECUTABLE)
        if nmap_path is None:
            raise Exception("Unable to locate {0} in path".format(self.NMAP_EXECUTABLE))

        # Simplistic target parsing is all we need, as nmap can understand network ranges and
        # hostnames
        target = self.configuration['target'].split('//')[-1]

        # Launch nmap
        self.spawn(nmap_path, self._generate_switches() + [target])

    def do_process_stdout(self, data):
        """
        Callback function whenever the executable (nmap) generates output on standard out.  Just buffer it internally
        and use it upon executable completion.

        Args:
            data: output from nmap; this should ostensibly in the form of an XML blob generated upon completion
        """
        self.stdout += data

    def do_process_stderr(self, data):
        """
        Callback function whenever the executable (nmap) generates output on standard error.  Just buffer it internally
        and use it upon executable completion.

        Args:
            data: output from nmap; this is string data that can be used to determine why nmap failed its scan
        """
        self.stderr += data

    def do_process_ended(self, status):
        """
        Once the plugin has finished running, parse the nmap output, generate the issues, and end execution.

        Args:
            status: exit code from the plugin runner's execution
        """

        # Plugin manually stopped
        if self.stopping and status == 9:
            self.report_finish("STOPPED")

        # Plugin completed execution normally
        elif status == 0:
            # Parse the XML output from nmap, sending in the baseline if it's available
            result = self.parse_nmap_xml(self.stdout,
                                         self.configuration.get('baseline', {}))

            # Based on the list of non-baseline open ports generated by parse_nmap_xml, create all the issues to pass
            # up to minion-backend
            issues = self.generate_issues(result)

            # Report the issues up to minion-backend
            self.report_issues(issues)

            # Callback to announce completion of execution
            self.report_finish()

        # If the plugin failed execution, stuff its stderr into a blob, along with the system's hostname to report
        # the failure up to minion-backend
        else:
            failure = {
                "hostname": socket.gethostname(),
                "exception": self.stderr,
                "message": "Plugin failed"
            }
            self.report_finish("FAILED", failure)
