#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
    hacked up from opencollab-import-mresolved
    @copyright: 2008-2009 Lari Huttunen, Juhani Eronen, Erno Kuusela
    @license: MIT <http://www.opensource.org/licenses/mit-license.php>
"""
import os
import sys
import time
import copy
import socket
import optparse
import cPickle as pickle
from opencollab.meta import Metas
from opencollab.util.wiki import importMetas
from opencollab.util.config import parseOptions
from opencollab.wiki import CLIWiki, WikiFailure


def parse_nvd_data(fname, metas, sw_keywords=''):
    new_metas = copy.deepcopy(metas)
    from lxml import objectify
    objectify.enable_recursive_str()
    tree = objectify.parse(open(fname))
    e = tree.getroot().entry
    do_nvd_entry(e, new_metas)
    while 1:
        e = e.getnext()
        if e is None:
            break
        do_nvd_entry(e, new_metas, sw_keywords)
    return new_metas


def do_nvd_entry(e, metas, sw_keywords=''):
    cveid = e.get('id')
    m = metas[cveid]

    for s in 'published-datetime', 'last-modified-datetime', 'security-protection', 'summary', 'severity':
        try:
            x = e['{http://scap.nist.gov/schema/vulnerability/0.4}' + s]
            m[s] = [unicode(x)]
            #print s, x.encode('ascii', 'replace')
        except AttributeError:
            pass

    try:
        vs = e['{http://scap.nist.gov/schema/vulnerability/0.4}vulnerable-software-list']
    except AttributeError:
        vs = []
    for s in vs:
        #print 'vulnerable-software', s['{http://scap.nist.gov/schema/vulnerability/0.4}product']
        m['vulnerable-software'] = [unicode(s['{http://scap.nist.gov/schema/vulnerability/0.4}product'])]

    try:
        refs = e['{http://scap.nist.gov/schema/vulnerability/0.4}references']
    except AttributeError:
        refs = []

    for r in refs:
        #print 'ref',
        #print r['{http://scap.nist.gov/schema/vulnerability/0.4}source'],
        #print r['{http://scap.nist.gov/schema/vulnerability/0.4}reference'].get('href')
        m['reference'] = [unicode(r['{http://scap.nist.gov/schema/vulnerability/0.4}source']) + u' ' + unicode(r['{http://scap.nist.gov/schema/vulnerability/0.4}reference'].get('href'))]

    metric = lambda x: e['{http://scap.nist.gov/schema/vulnerability/0.4}cvss']['{http://scap.nist.gov/schema/cvss-v2/0.2}base_metrics']['{http://scap.nist.gov/schema/cvss-v2/0.2}' + x]

    for mn in 'score', 'access-vector', 'access-complexity', 'availability-impact', 'confidentiality-impact', 'integrity-impact':
        try:
            z = metric(mn)
            #print 'cvss-'+mn, z
        except AttributeError:
            pass
        else:
            m['cvss-' + mn] = [unicode(z)]

    try:
        cweid = e['{http://scap.nist.gov/schema/vulnerability/0.4}cwe'].get("id")
    except AttributeError:
        pass
    else:
        if cweid:
            m["cwe-id"] = [cweid]

    m['gwikicategory'] = ['CategoryVulnerability']

    if sw_keywords:
        # only keep vulns that have keyword matches in "vulnerable-software"
        ok = 0
        for w in sw_keywords.split():
            for s in m['vulnerable-software']:
                if w in s:
                    ok = 1
        if not ok:
            del metas[cveid]
        #else:
        #    print 'leave', cveid, s


def main():
    parser = optparse.OptionParser()
    parser.add_option("-k", "--sw-keywords",
                      dest="swkeywords",
                      default='',
                      metavar="KEYWORDS",
                      help="Only include entries with one of KEYWORDS in the 'vulnerable-software' field")
    parser.add_option("-p", "--pickle",
                      dest="picklefn",
                      default=None,
                      metavar="PICKLE-FILENAME",
                      help="PICKLE-FILENAME to use instead of wiki")
    parser.set_usage("%prog [options] 1..N input files")
    failed = []
    sect = "import-nvd-xml"
    ops = parseOptions(parser, sect, template=True)
    url = ops['creds']['url']
    x509 = ops[sect]["x509"]
    x509_ca_file = ops[sect]["x509_ca_file"]
    verbose = ops[sect]['verbose']
    template = ops[sect]["template"]
    if template is None:
        template = "CveTemplate"
    sw_keywords = ops[sect]["swkeywords"]
    picklefn = ops[sect]["picklefn"]
    args = ops[sect]["args"]
    if picklefn and os.path.exists(picklefn):
        start = time.time()
        metas = pickle.load(open(picklefn))
        if verbose:
            print 'NOTE: Loaded %d metas from %s in %.2f s' % (len(metas), picklefn, time.time() - start)
    else:
        metas = Metas()
    if len(args) < 1:
        parser.error("You need to specify 1..N input files. Use -h for help.")
    else:
        for file in args:
            metas.update(parse_nvd_data(file, metas, sw_keywords))
    if picklefn and os.path.exists(picklefn):
        if verbose:
            print 'NOTE: Writing metas to', picklefn
        start = time.time()
        f = open(picklefn, 'w')
        pickle.dump(metas, f, 2)
        f.close()
        if verbose:
            print 'NOTE: done (%.2f s)' % (time.time() - start)
    else:
        while True:
            try:
                collab = CLIWiki(ssl_verify_cert=x509, ssl_ca_certs=x509_ca_file, **ops['creds'])
            except WikiFailure:
                print "ERROR: Authentication failed."
            except (UnicodeError, socket.gaierror):
                sys.exit("ERROR: Not a valid URL.")
            else:
                break
        if verbose:
            print "NOTE: Importing metas to", url
        failed = importMetas(collab, metas, template, verbose, False)
        if failed:
            for page in failed:
                if verbose:
                    print "ERROR: Uploading",  page, "metas", repr(metas[page]), "failed."
                else:
                    print "ERROR: Uploading:", page, "metas failed."

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print "Script interrupted via CTRL-C."
