#! /usr/bin/python
# Copyright (C) 2008 by Intevation GmbH
# Authors:
# Bernhard Herzog <bh@intevation.de>
#
# This program is free software under the GPL (>=v2)
# Read the file COPYING coming with the software for details.


"""Script to build a mapping of OpenVAS NVT (o)ids to filenames.
Run it with the --help option for information about the option.

A typical example for OpenVAS:
  build_oid_map.py --workdir=/tmp/openvas-oid-mapping \
     --svn-url=https://svn.wald.intevation.org/svn/openvas/trunk/openvas-plugins/scripts \
     --output=oidmapping.txt

With the --rsync-to option it can also rsync the file to a remote system:

     --rsync-to=wald.intevation.org:/openvas/htdocs/oidmapping.txt

The mapping generated by this script is used by the openvas Website to
redirect URLs of the form http://www.openvas.org/?oid=<OID> to a page
with information about the specified plugin.
"""


import sys
import os
import tempfile
import subprocess
import errno
import re
from optparse import OptionParser

# OID prefix for the legacy script ids
openvas_oid_prefix = "1.3.6.1.4.1.25623.1.0."


def ensure_directory(directory):
    """Creates directory and all its parents.

    Unlike os.makedirs, this function doesn't throw an exception if the
    directory already exists.
    """
    if not os.path.isdir(directory):
        os.makedirs(directory)


class SubprocessError(EnvironmentError):

    def __init__(self, command, returncode):
        EnvironmentError.__init__(self,
                                  "Command %r finished with return code %d"
                                  % (command, returncode))
        self.returncode = returncode


def call(command, suppress_output=False, **kw):
    """Run command as a subprocess and wait until it is finished.

    The command should be given as a list of strings to avoid problems
    with shell quoting.  If the command exits with a return code other
    than 0, a SubprocessError is raised.
    """
    if suppress_output:
        kw["stdout"] = open(os.devnull, "w")
        kw["stderr"] = open(os.devnull, "w")
    ret = subprocess.call(command, **kw)
    if ret != 0:
        raise SubprocessError(command, ret)


def svn_checkout(url, localdir, verbose):
    """Runs svn to checkout the repository at url into the localdir"""
    cmd = ["svn", "checkout"]
    if verbose == 0:
        cmd.append("-q")
    cmd.extend([url, localdir])
    call(cmd)

def svn_update(localdir, verbose):
    """Runs svn update on the localdir."""
    cmd = ["svn", "update"]
    if verbose == 0:
        cmd.append("-q")
    cmd.extend([localdir])
    call(cmd)

def svn_checkout_or_update(url, localdir, verbose=1):
    """Updates the working copy of url in localdir.

    If localdir doesn't exist yet, url is checked out into that
    directory.
    """
    if os.path.exists(localdir):
        if verbose:
            print >>sys.stderr, "svn update in %s" % (localdir,)
        svn_update(localdir, verbose)
    else:
        if verbose:
            print >>sys.stderr, "checking out %s into %s" % (url, localdir)
        svn_checkout(url, localdir, verbose)


class OIDMapBuilder(object):

    def __init__(self, basedir, verbose):
        self.basedir = basedir
        self.mapping = dict()
        self.verbose = verbose

    def scan(self):
        for filename in os.listdir(self.basedir):
            if filename.endswith(".nasl"):
                self.scan_file(filename)

    def scan_file(self, filename):
        for line in open(os.path.join(self.basedir, filename)):
            line = line.strip()
            if line.startswith("#"):
                # skip comments
                continue
            match = re.search(r"script_id\s*\(\s*([0-9]+)\s*\)", line)
            if match:
                oid = openvas_oid_prefix + match.group(1)
                if not oid in self.mapping:
                    self.mapping[oid] = filename
                else:
                    if self.verbose > 0:
                        print >>sys.stderr, ("'%s' and '%s' have the same id"
                                             % (filename, self.mapping[oid]))
                break
        else:
            if self.verbose > 0:
                print>>sys.stderr, "Could not find script_id in %s" % filename

    def write_mapping(self, filename):
        """Writes the mapping to the file named by filename"""
        directory = os.path.dirname(filename)
        fileno, tempname = tempfile.mkstemp(".txt", "oidmapping", directory)
        try:
            outfile = os.fdopen(fileno, "w")
            for key, value in self.mapping.iteritems():
                outfile.write("%s %s\n" % (key, value))
            outfile.flush()
            os.rename(tempname, filename)
            outfile.close()
        finally:
            try:
                os.remove(tempname)
            except OSError, exc:
                if exc.errno == errno.ENOENT:
                    # should only happen if tempname has already been
                    # renamed and therefore doesn't exist anymore under that
                    # name
                    pass
                else:
                    raise

def rsync_mapping(filename, remote_filename, verbose):
    """Runs rsync to copy filename to remote_filename."""
    # ignore the verbose argument for now, because rsync is quiet by
    # default
    call(["rsync", filename, remote_filename])


def main():
    parser = OptionParser()
    parser.set_defaults(verbose=1, svn_update=True)
    parser.add_option("-o", "--output", help="The output file")
    parser.add_option("--workdir", help=("Working directory"
                                         " (holds svn checkout and log files)"))
    parser.add_option("--svn-url",
                      help=("URL to use for initial checkout into"
                            " $workdir/checkout"))
    parser.add_option("--rsync-to",
                      help=("rsync parameter for the remote filename to"
                            " rsync the mapping to"))
    parser.add_option("-q", "--quiet",
                      action="store_const", const=0, dest="verbose",
                      help="Do not print diagnostic output to stderr")
    parser.add_option("--no-update",
                      action="store_const", const=False, dest="svn_update",
                      help="Do not update the working copy.")


    options, args = parser.parse_args()

    ensure_directory(options.workdir)

    checkout_dir = os.path.join(options.workdir, "checkout")

    if options.svn_update:
        svn_checkout_or_update(options.svn_url, checkout_dir, options.verbose)

    builder = OIDMapBuilder(checkout_dir, options.verbose)
    builder.scan()

    if options.output:
        builder.write_mapping(options.output)

    if options.verbose > 0:
        print >>sys.stderr, "total: %d files" % len(builder.mapping)

    if options.rsync_to:
        if options.verbose > 0:
            print >>sys.stderr, "rsyncing to %s" % options.rsync_to
        rsync_mapping(options.output, options.rsync_to, options.verbose)

main()
