#!/usr/bin/env python

# OpenVAS: joinreport.py:
# $Id$
# Description:
# Simple script to join multiple OpenVAS Reports into a single
# preserving host details.
# Current version is a prototype helper script and not save against
# multiple identical IPs.
#
# Authors:
# Hani Benhabilis <hani.benhabilis@greenbone.net>
#
# Copyright:
# Copyright (C) 2013 Greenbone Networks GmbH
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
#/

import sys
import uuid
from xml.dom import minidom

def new_uuid():
    return str(uuid.uuid4())


def create_report(all_results, all_hosts, all_host_start, all_host_end):
    # Create new report
    new_doc = minidom.Document()
    report = new_doc.createElement('report')
    new_doc.appendChild(report)
    # Create inner report
    inreport = new_doc.createElement('report')
    report.appendChild(inreport)
    report.setAttribute('id', new_uuid())
    report.setAttribute('format_id', 'a994b278-1f62-11e1-96ac-406186ea4fc5')
    report.setAttribute('extension', 'xml')
    report.setAttribute('content_type', 'text/xml')
    report.setAttribute('type', 'scan')

    # Now, we work within the "inner" report

    # Add results
    new_results = new_doc.createElement('results')
    inreport.appendChild(new_results)
    for result in all_results:
        new_results.appendChild(result)

    # <host_start> and <host_end>
    for host_start in all_host_starts:
        inreport.appendChild(host_start)
    for host_end in all_host_ends:
        inreport.appendChild(host_end)


    # Add host (details)
    for _, host in all_hosts.items():
        new_host = new_doc.createElement('host')
        new_host.appendChild(host['ip'])
        new_host.appendChild(host['start'])
        new_host.appendChild(host['end'])
        for detail in host['details']:
            new_host.appendChild(detail)
        inreport.appendChild(new_host)

    return new_doc.toxml(encoding='UTF-8')

def create_report_node(xml_file):
    """ Gets the inner <report> element from an exported OpenVAS report.xml """

    xmldoc = minidom.parse(xml_file)
    # Get the "inner" report.
    for node in xmldoc.childNodes:
        if node.localName == "report":
            reportdoc = node
    for node in reportdoc.childNodes:
        if node.localName == "report":
            reportdoc = node
            break
    return reportdoc

def get_report_results(reportdoc):
    """ Gets a list of results from a <report> element. """

    # Handle the "results" part:
    # Merge all "result" elements into all_results.
    for node in reportdoc.childNodes:
        if node.localName == "results":
            resultsdoc = node
    return [r for r in resultsdoc.childNodes if r.localName == 'result']

def get_report_nodes(reportdoc, node_name):
    """ Gets a list of child nodes with provided name from a <report> element. """

    return [r for r in reportdoc.childNodes if r.localName == node_name]

def fill_report_hosts(reportdoc, all_hosts):
    """ Fills <host> elements in all_hosts table. If a host with the same IP
    exists in the table, host details of the new element will be merged into the
    same element.
    """

    # Handle the "host" elements (host details) part:
    # Creating new table entry for every host (by IP).
    hosts = [h for h in reportdoc.childNodes if h.localName == 'host']
    for host in hosts:
        ipnode = None
        start_node = None
        end_node = None
        details_nodes = list()
        for e in host.childNodes:
            if e.localName == 'ip':
                ipnode = e
            elif e.localName == 'start':
                start_node = e
            elif e.localName == 'end':
                end_node = e
            elif e.localName == 'detail':
                details_nodes.append(e)
            elif e.localName == None and e.nodeType == e.TEXT_NODE:
                # <host>'s text, which is empty.
                pass
            #else:
            #    assert False, "<host> has element {0} other than ip/start/end/detail".format(e.localName)

        hostip = ipnode.childNodes[0].nodeValue
        if all_hosts.get(hostip) is None:
            all_hosts[hostip] = dict()
            all_hosts[hostip]['ip'] = ipnode
            all_hosts[hostip]['start'] = start_node
            all_hosts[hostip]['end'] = end_node
            all_hosts[hostip]['details'] = details_nodes
        else:
            sys.stderr.write("Host details duplication for {0}".format(hostip))

if __name__ == '__main__':

    # <result> nodes list from all reports.
    all_results = list()
    # <host> nodes hash table with ip addresses as key
    all_hosts = dict()
    # <host_start> nodes lists.
    all_host_starts = list()
    # <host_end> nodes lists.
    all_host_ends = list()

    for xml_file in sys.argv[1:]:
        reportdoc = create_report_node(xml_file)
        all_results.extend(get_report_results(reportdoc))
        fill_report_hosts(reportdoc, all_hosts)
        all_host_starts.extend(get_report_nodes(reportdoc, 'host_start'))
        all_host_ends.extend(get_report_nodes(reportdoc, 'host_end'))

    #print "Len results", len(all_results)

    new_report_xml = create_report(all_results, all_hosts, all_host_starts,
                                   all_host_ends)
    sys.stdout.write(new_report_xml)

