<commit_before>import xml.etree.ElementTree
import logging

import database
import errors


class Import():

    def __init__(self, db_file):
        self.log = logging.getLogger('IMPORT')
        self.db = database.ScanDatabase(db_file)

    def import_scan(self, scan_data):
        file_type = self.get_file_type(scan_data[:100])

        try:
            if file_type == 'Nessus':
                self.import_nessus(scan_data)
            elif file_type == 'Nmap':
                self.import_nmap(scan_data)
            else:
                self.log.error('Unknown file type, skipping import.')
                return 'Failed'

        except (errors.ScanImportError):
            self.log.error('Failed to parse scan file.')
            return 'Failed'

        return 'Succeeded'

    def get_file_type(self, header):
        """
        Determine the source of the scan data.
        """
        self.log.debug('Checking file type with header {0}.'.format(header))
        if '<NessusClientData_v2>' in header:
            return 'Nessus'
        elif '<!DOCTYPE nmaprun>' in header:
            return 'Nmap'
        else:
            return 'Unknown'

    def import_nessus(self, scan_data):
        """
        Load the Nessus scan data into an XML structure and import the data.
        """
        self.log.info('Importing Nessus file.')
        # Load Nessus XML file into the tree and get the root element.
        try:
            root = xml.etree.ElementTree.fromstring(scan_data)

            for report in root.findall('Report'):
                report_name = report.attrib['name']
                self.log.info('Importing report {0}.'.format(report_name))
                self.process_nessus_hosts(report.findall('ReportHost'))

        except xml.etree.ElementTree.ParseError:
            self.log.error('Unable to parse Nessus XML file.')
            raise errors.ScanImportError

    def process_nessus_hosts(self, report_hosts):
        """
        Process each host in a report
        """
        self.log.debug('Processing Nessus report hosts.')
        for host in report_hosts:
            self.log.debug('Getting IP address for report host.')
            ip = ''
            for tag in host.find('HostProperties').findall('tag'):
                if tag.attrib['name'] == 'host-ip':
                    ip = tag.text

            if (ip != ''):
                self.log.info('Processing ip {0}.'.format(ip))
                self.process_nessus_items(ip, host.findall('ReportItem'))

    def process_nessus_items(self, ip, report_items):
        """
        Process each report item in a host.
        """
        for item in report_items:
            text = xml.etree.ElementTree.tostring(item, encoding='utf-8')
            self.log.debug('Processing report item {0}.'.format(text))

            port = int(item.attrib['port'])
            proto = item.attrib['protocol']

            note = '--{0}--\n\n'.format(item.attrib['pluginID'])
            description = item.find('description')
            output = item.find('plugin_output')

            if description is not None:
                note += 'Description:\n{0}'.format(description.text.strip('\n'))

            if output is not None:
                note += 'Output:\n{0}'.format(output.text.strip('\n'))

            if self.db.create_item(ip, port, proto, note) is False:
                self.log.error('Unable to create new Nessus item in database.')

    def import_nmap(self, scan_data):
        """
        Load the Nmap scan data into an XML structure and import the data.
        """
        self.log.info('Importing Nmap file.')
        # Load Nmap XML file into the tree and get the root element.
        try:
            root = xml.etree.ElementTree.fromstring(scan_data)

            for host in root.findall('host'):
                self.log.debug('Getting IP address for host.')
                address = host.find('address')
                ip = address.attrib['addr']

                self.log.info('Processing ip {0}.'.format(ip))
                self.process_nmap_ports(ip, host.findall('ports/port'))
                self.process_nmap_hostscripts(ip, host.findall('hostscript/script'))<commit_msg>Document any available Metasploit modules in item note.<commit_after>import xml.etree.ElementTree
import logging

import database
import errors


class Import():

    def __init__(self, db_file):
        self.log = logging.getLogger('IMPORT')
        self.db = database.ScanDatabase(db_file)

    def import_scan(self, scan_data):
        file_type = self.get_file_type(scan_data[:100])

        try:
            if file_type == 'Nessus':
                self.import_nessus(scan_data)
            elif file_type == 'Nmap':
                self.import_nmap(scan_data)
            else:
                self.log.error('Unknown file type, skipping import.')
                return 'Failed'

        except (errors.ScanImportError):
            self.log.error('Failed to parse scan file.')
            return 'Failed'

        return 'Succeeded'

    def get_file_type(self, header):
        """
        Determine the source of the scan data.
        """
        self.log.debug('Checking file type with header {0}.'.format(header))
        if '<NessusClientData_v2>' in header:
            return 'Nessus'
        elif '<!DOCTYPE nmaprun>' in header:
            return 'Nmap'
        else:
            return 'Unknown'

    def import_nessus(self, scan_data):
        """
        Load the Nessus scan data into an XML structure and import the data.
        """
        self.log.info('Importing Nessus file.')
        # Load Nessus XML file into the tree and get the root element.
        try:
            root = xml.etree.ElementTree.fromstring(scan_data)

            for report in root.findall('Report'):
                report_name = report.attrib['name']
                self.log.info('Importing report {0}.'.format(report_name))
                self.process_nessus_hosts(report.findall('ReportHost'))

        except xml.etree.ElementTree.ParseError:
            self.log.error('Unable to parse Nessus XML file.')
            raise errors.ScanImportError

    def process_nessus_hosts(self, report_hosts):
        """
        Process each host in a report
        """
        self.log.debug('Processing Nessus report hosts.')
        for host in report_hosts:
            self.log.debug('Getting IP address for report host.')
            ip = ''
            for tag in host.find('HostProperties').findall('tag'):
                if tag.attrib['name'] == 'host-ip':
                    ip = tag.text

            if (ip != ''):
                self.log.info('Processing ip {0}.'.format(ip))
                self.process_nessus_items(ip, host.findall('ReportItem'))

    def process_nessus_items(self, ip, report_items):
        """
        Process each report item in a host.
        """
        for item in report_items:
            text = xml.etree.ElementTree.tostring(item, encoding='utf-8')
            self.log.debug('Processing report item {0}.'.format(text))

            port = int(item.attrib['port'])
            proto = item.attrib['protocol']

            note = '--{0}--\n\n'.format(item.attrib['pluginID'])
            description = item.find('description')
            output = item.find('plugin_output')
            metasploit = item.find('metasploit_name')

            if description is not None:
                note += 'Description:\n{0}'.format(description.text.strip('\n'))

            if output is not None:
                note += 'Output:\n{0}'.format(output.text.strip('\n'))

            if metasploit is not None:
                note += 'Metasploit: {0}\n'.format(metasploit.text)

            if self.db.create_item(ip, port, proto, note) is False:
                self.log.error('Unable to create new Nessus item in database.')

    def import_nmap(self, scan_data):
        """
        Load the Nmap scan data into an XML structure and import the data.
        """
        self.log.info('Importing Nmap file.')
        # Load Nmap XML file into the tree and get the root element.
        try:
            root = xml.etree.ElementTree.fromstring(scan_data)

            for host in root.findall('host'):
                self.log.debug('Getting IP address for host.')
                address = host.find('address')
                ip = address.attrib['addr']

                self.log.info('Processing ip {0}.'.format(ip))
                self.process_nmap_ports(ip, host.findall('ports/port'))
                self.process_nmap_hostscripts(ip, host.findall('hostscript/script'))
{'fraction_non_alphanumeric': 0.11554524361948956, 'fraction_numerical': 0.0031322505800464038, 'mean_word_length': 1.9303195105370496, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 10}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5728385', 'n_tokens_mistral': 2490, 'n_tokens_neox': 2385, 'n_words': 598}