#!/usr/bin/python
# Copyright 1999-2008 Gentoo Foundation
# Distributed under the terms of the GNU General Public License v2

# Written by Robert Buchholz <rbu@gentoo.org> and Tobias Heinlein <keytoaster@gentoo.org>
# Based on a Perl script by Stefan Fritsch <sf@sfritsch.de>

# This program needs app-portage/eix installed and its database up-to-date
# It uses a modified NVD python component

import string
import sys
import os
import re
import subprocess

DEFAULT_ISSUE_REGEX = '^CVE-20(0[3-9]|10)'
DEFAULT_TODO_REGEX = '^\s+TODO: check$'


def filterstring(strng):
    """ Make a string translation filter that converts all illegal chars to spaces
    """
    import string

    strng = strng.encode('ascii', 'ignore')
    allowed = string.letters + string.digits + '_'

    allchars = string.maketrans(u'', u'') # table of all chars
    replacechars = ''.join([c not in allowed and ' ' or c for c in allchars])

    return strng.translate(replacechars).encode('utf-8')


class EntryEditor:
    def __init__(self, issue_regex, todo_regex, replace_line, list_only, list_lines, glsa_style, sort_only, bugreporter):
        self.datafile = "./data/CVE/list"
        self.issue_regex = issue_regex
        self.todo_regex = todo_regex
        self.replace_line = replace_line
        self.editor = os.environ.get('EDITOR', os.environ.get('VISUAL', "nano"))
	self.browser = os.environ.get('BROWSER')

        self.listdata = self.read_list(self.datafile)

        if sort_only:
            self.save()
            return

        self.cvedb = CVEData()

        todos = self.filter_todos(self.listdata)

        # list of entries to post a bug for
        self.bugs_collected = {}
        self.bugs_atom = []
        self.bugreporter = bugreporter

        if list_only:
            self.list_entries(todos, list_lines, glsa_style)
        else:
            self.bugupdates = {}
            self.recently_saved = True
            self.work_entries(todos)


    def read_list(self, filename):
        file = open(filename)
        entries = []
        cur_entry = []

        for line in file:
            if line[0:3] == 'CVE' and cur_entry:
                entries.append(cur_entry)
                cur_entry = []
            cur_entry.append(line)
        if cur_entry:
            entries.append(cur_entry)
        file.close()
        entries.sort()
        return entries


    def filter_todos(self, entries):
        todo_matcher = re.compile(self.todo_regex)
        issue_matcher = re.compile(self.issue_regex)
        todos = []

        for entry in entries:
            if not issue_matcher.search(entry[0]):
                continue
            for line in entry:
                if todo_matcher.search(line):
                    todos.append(entry)
                    break
        return todos


    def list_entries(self, entries, max_line_count, glsa):
        cvelist = []
        if glsa:
            print "<ul>"
        for entry in entries:
            cve = self.cvedb.get_cve_from(entry)
            if not cve:
                continue
            cvelist.append(cve)
            if glsa:
                names = []
                for cvenames in self.cvedb.guess_name_for(cve).values():
                    names.extend(cvenames)
                joinednames = u", ".join(list(set(names)))
                names = joinednames
                if names:
                    credit = names + " reported "
                else:
                    credit = ""
                print "<li>"
                cvedesc = self.cvedb.get_cve_desc(cve, indentation = 0)
                if cvedesc[-1] == '.':
                    cvedesc = cvedesc[:-1]
                print u"%s%s (%s)." % (credit, cvedesc, cve)
                print "</li>"
            else:
                print "%s (http://nvd.nist.gov/nvd.cfm?cvename=%s):" % (cve, cve)
                cvedesc = self.cvedb.get_cve_desc(cve, indentation = 2)
                cvedesc = '\n'.join(cvedesc.split('\n')[:int(max_line_count)])
                print cvedesc
        if glsa:
            print "</ul>"
        else:
            print "\nAddressed %d CVEs: %s" % (len(cvelist), self.bugs_unify_cvenames(cvelist))


    def work_entries(self, entries):
        self.entrynum = 0
        while self.entrynum >= 0 and self.entrynum < len(entries):
            entry = entries[self.entrynum]
            result = self.cvedb.print_all_about(entry)
            if not result:
                self.entrynum += 1
                continue
            print "CVE # in list: " + str(self.entrynum + 1) + " / " + str(len(entries))
            cve = result[0]
            product = filterstring(result[1])

            # By default, do an auto-search
            command = 'p'
            while command != '\n':
                product = self.handlecommand(command, entry, cve, product)
                for line in entry:
                    print line,
                print
                print ">>> ",
                command = sys.stdin.readline()
            self.entrynum += 1
        print
        if not self.recently_saved:
             print "All entries done, attempting to quit, do you want to save? (Y/n) ",
             want_to_save = sys.stdin.readline()
             if want_to_save[0] == "y" or want_to_save[0] == "\n":
                 self.save()
        self.update_bugs()


    def update_bugs(self):
        if len(self.bugupdates) > 0:
             print "I have %d bug updates for you ..." % len(self.bugupdates)
             for bug in self.bugupdates.items():
                 self.update_bug(*bug)
        self.bugupdates = {}


    def update_bug(self, bug, cves):
        title = self.bugreporter.get_bug_title(bug)
        if not title:
            return
        bug_cves = self.bugreporter.get_bug_cves(bug, title = title)
        total_cves = list(set(bug_cves + cves))
        new_title = "%s (%s)" % (BugReporter.CVEGROUPALL.sub('', re.sub('\(?CVEs? requested\)?', '', title)).strip(),
                          self.bugs_unify_cvenames(total_cves))
        description = ""
        for cve in cves:
            description += "%s (http://nvd.nist.gov/nvd.cfm?cvename=%s):\n" % (cve, cve)
            description += "%s" % self.cvedb.get_cve_desc(cve, indentation = 2)
            description += "\n\n"

        print "=========== Updating bug %s ===========" % (bug)
        print " Title: %s" % (title)
        print "     => %s" % (new_title if title != new_title else "(unchanged)")
        print " Comment: "
        print description

        print "Commit changes to BZ? [Y/n]: ",

        answer = sys.stdin.readline()
        if answer[0] == "y" or answer[0] == "Y" or answer[0] == "\n":
            self.bugreporter.modify_bug(bug, new_title, description)



    def handlecommand(self, command, entry, cve, product):
        if not command:
            command = "q"
        if command[0] == 'e' or command[0] == 'v':
            self.edit_entry(entry)
        
	elif command[0] == 's':
            self.save()

        elif command[0] == 'c':
            self.bugs_collected[cve] = entry
            self.bugs_atom = product
        
        elif command[0] == 'B':
            self.open_in_browser(cve)
        
	elif command[0] == 'b':
            self.file_bug(entry, cve, product)
        
        elif command[0] == 'q':
            if self.recently_saved == False:
                self.save()
            sys.exit()
        
        elif command[0] == 'p' or command[0] == 'd':
            if len(command) > 2:
                product = command[1:].strip()
            if product:
                print
                res = self.get_best_results(product)
                if not res:
                    print "Couldn't find a hint of this. (%s)" % (product.strip())
                elif len(res) == 1:
                    print "Found the package in our tree (%s):" % (res[0].strip())
                    os.spawnlp(os.P_WAIT, 'eix', 'eix', '--compact', res[0].strip())
                elif len(res) < 11 or command[0] == 'd':
                    res = [atom.strip() for atom in res]
                    print "Found multiple packages in our tree (%s): " % (product.strip())
                    regexstring = '|'.join(res)
                    os.spawnlp(os.P_WAIT, 'eix', 'eix', '--compact', regexstring)
                else:
                    print "Found %d packages in our tree (%s). Press d to display." % (len(res), product.strip())

        elif command[0] == 'P':
            if len(command) > 2:
                product = command[1:].strip()
            if product:
                from subprocess import Popen, PIPE
                myproc = Popen(["bugz","search","--quiet","--columns=1000", "--product=Gentoo Security", product], stdout=PIPE)
                myproc.wait
                res = myproc.communicate()[0]
                res = str.splitlines(res)
                if not res:
                    print "No results from our bugzilla. (%s)" % (product.strip())
                else:
                    for i in range (len(res)):
                        print "http://bugs.gentoo.org/" + res[i]
				
        elif command == '\n':
            # next one, please
            pass

        elif command[0] == 'n':
            # The whole string is only "n": use the product as nfu
            if len(command.strip()) == 1:
                notforus = product
            else:
                notforus = command[1:]
            self.update_entry_todo(entry, "\tNOT-FOR-US: %s" % (notforus.strip()))

        elif command[0] == '#':
            bugno = command[1:].strip()
            self.update_entry_todo(entry, "\tBUG: %s" % (bugno))
            self.bugupdates.setdefault(bugno, []).append(cve)
        
        elif command[0] == 'l':
            self.update_entry_todo(entry, "\tTODO: check-later")       

        elif command[0] == 'u':
            self.update_bugs()

        elif command[0] == '<':
            if len(command) > 2:
                self.entrynum -= int(command[1]) + 1
            else:
                self.entrynum -= 1 + 1

        elif command[0] == '>':
            if len(command) > 2:
                self.entrynum += int(command[1]) - 1
            else:
                self.entrynum += 1 - 1

        else:
            print
            print "Available actions:"
            miniusage()

        return product

    def update_entry_todo(self, entry, replacement):
            replacement = "%s\n" % (replacement.rstrip())
            if self.replace_line:
                matcher = re.compile(self.todo_regex)
                for idx, line in enumerate(entry):
                    if matcher.match(line):
                        entry[idx] = replacement
            else:
                entry.append(replacement)
            self.recently_saved = False


    def edit_entry(self, entry):
        import tempfile
        (fd, filename) = tempfile.mkstemp(suffix='.txt', prefix='gsec', text=True)
        for line in entry:
            os.write(fd, line)
        os.close(fd)
        os.spawnlp(os.P_WAIT, self.editor, self.editor, filename)

        # This edits our reference to the list, and therefore also the master entry
        del entry[:]
        newentry = open(filename)
        entry.extend(newentry.readlines())
        newentry.close()
        os.remove(filename)
        self.recently_saved = False


    def open_in_browser(self, cve):
	url = "http://nvd.nist.gov/nvd.cfm?cvename=%s" % (cve)
	if not self.browser:
	    print "Please set the environment variable BROWSER to use this feature."
	else:
	    os.spawnlp(os.P_WAIT, self.browser, self.browser, url)


    def file_bug(self, entry, cve, product):
        cvelist = self.bugs_collected.keys()
        cvelist.sort()

        if len(cvelist) < 1:
            # No bugs collected yet. Add the current one.
            self.bugs_collected[cve] = entry
            self.bugs_atom = product
            cvelist = self.bugs_collected.keys()

        res = self.get_best_results(self.bugs_atom)
        if res and len(res) == 1:
            self.bugs_atom = res[0].strip()

        title = self.bugs_atom + " DESCR (" + self.bugs_unify_cvenames(cvelist) + ")"
        description = ""

        bug_url = ""
        for cve in cvelist:
            # print cve
            description += "%s (http://nvd.nist.gov/nvd.cfm?cvename=%s):\n" % (cve, cve)
            description += "%s" % self.cvedb.get_cve_desc(cve, indentation = 2)
            description += "\n\n"
            if not bug_url:
                for source, url in self.cvedb.get_refs_for(cve):
                    if source == u"CONFIRM":
                        bug_url = url
                        break

        import tempfile
        (fd, filename) = tempfile.mkstemp(suffix='.txt', prefix='gsec', text=True)
        os.write(fd, "%s\n" % (title))
        os.write(fd, "\n== Only the first line is used as a title, comments follow ==\n")
        os.write(fd, description)
        os.close(fd)
        os.spawnlp(os.P_WAIT, self.editor, self.editor, filename)

        editfile = open(filename)
        title = editfile.readline()
        editfile.close()
        os.remove(filename)

        print "Title: %s" % (title)
        print "Description:"
        print description
        print
        print "Component for bug [Enter=Vuln, k=Kernel, otherwise free text]: ",
        comp = sys.stdin.readline()
        component = ""
        if comp[0] == "a":
            component = "Auditing"
        elif comp[0] == "k":
            component = "Kernel"
        elif comp[0] == "\n":
            component = "Vulnerabilities"
        else:
            component = comp.strip()
        print "Initial Whiteboard: ",
        whiteboard = sys.stdin.readline()

        print "Will file this bug [Y/n]: ",

        answer = sys.stdin.readline()
        if answer[0] == "y" or answer[0] == "Y" or answer[0] == "\n":
            bugno = self.bugreporter.post_bug(title, description, component, whiteboard, bug_url)
            if bugno:
                print "Success! https://bugs.gentoo.org/%d" % (bugno)
                # note the bug in the list
                for cve in self.bugs_collected.keys():
                    entry = self.bugs_collected[cve]
                    self.update_entry_todo(entry, "\tBUG: %d" % (bugno))
            else:
                print "Failed creating the bug"

        ## This edits our reference to the list, and therefore also the master entry
        #del entry[:]
        self.bugs_atom = ""
        self.bugs_collected = {}


    def bugs_unify_cvenames(self, cvelist):
        """ This will make a list of CVE names readable for bugzilla.
        CVE-2007-1234 CVE-2007-1235 CVE-2007-1236 -> CVE-2007-{1234,1235,1236}
        """
        matcher = re.compile("CVE-(\d+)-(\d+)")

        cveyears = {}
        for cvename in cvelist:
            match = matcher.match(cvename)
            year = match.group(1)
            id = match.group(2)
            cveyears.setdefault(year, []).append(id)

        title = ""
        cvekeys_sorted = cveyears.keys()
        cvekeys_sorted.sort()
        for year in cvekeys_sorted:
            title += "CVE-%s" % (year)
            if len(cveyears[year]) == 1:
                title += "-%s," % (cveyears[year][0])
            else:
                title += "-{"
                cveyears[year].sort()
                for id in cveyears[year]:
                    title += "%s," % (id)
                title = title[:-1] + "},"
        # remove last comma
        title = title[:-1]
        return title



    def get_best_results(self, query):
        ''' Do not call with unfiltered strings '''
        
        searchresults = subprocess.Popen(['eix', '--only-names'] + query.split(), stdout=subprocess.PIPE).stdout.readlines()

        if searchresults and len(searchresults) > 0:
            return searchresults
        words = query.split()
        words.reverse()
        for word in words:
            searchresults = subprocess.Popen(['eix', '--only-names', word], stdout=subprocess.PIPE).stdout.readlines()
            if searchresults and len(searchresults) > 0 and len(searchresults) < 20:
                return searchresults

        searchresults = subprocess.Popen(['eix', '--only-names', '-S'] + query.split(), stdout=subprocess.PIPE).stdout.readlines()

        if searchresults and len(searchresults) > 0:
            return searchresults
        words = query.split()
        words.reverse()
        for word in words:
            searchresults = subprocess.Popen(['eix', '--only-names', '-S', word], stdout=subprocess.PIPE).stdout.readlines()
            if searchresults and len(searchresults) > 0 and len(searchresults) < 20:
                return searchresults

        return None

    def save(self):
        file = open(self.datafile, 'w')
        for entry in self.listdata:
            for line in entry:
                file.write(line)
        file.close()
        print
        print "Save completed."
        print
        self.recently_saved = True


def setup_paths():
    """ Set up paths to include our local lib dir """
    import os.path
    check_file = 'lib/python/nvd.py'
    path = os.getcwd()
    while 1:
        if os.path.exists("%s/%s" % (path, check_file)):
            sys.path = [path + '/lib/python'] + sys.path
            return path
        idx = string.rfind(path, '/')
        if idx == -1:
            raise ImportError, "could not setup paths"
        path = path[0:idx]


def miniusage():
    """ Print available actions """
    print ''' * RETURN    Goes to the next entry

 * B        open CVE in browser
 * pNAME    Does a new search in the tree for NAME
 * PNAME    Does a new search in bugzilla for NAME
 * e or v   Calls your $EDITOR to edit this entry
 * d        Redisplays the last search completely
 * n        Marks "NOT-FOR-US: NAME" while NAME is the last product (from p or guess)
 * nSTRING  Marks "NOT-FOR-US: STRING"
 * #123     Marks "BUG: 123"
 * u        Invokes bug update mechanism immediately
 * <[n]     Jumps back one [or n] entries in the list
 * >[n]     Jumps forward one [or n] entries in the list
 * Entering any other string displays this help

 * CTRL+C   Quits without saving
 * q        Quits with saving
 * s        Saves the current state without quitting
'''

def usage(programname):
    """ Print usage information """
    print "Usage: %s [-h] [-l [-n <X>]] [-i <regex>] [-t <regex>] [-T] [-a]" % (programname)
    print '''
This script reads entries from data/CVE/list and prints all items marked "TODO: check".

Parameters:
 -h        Display this help
 -l        Only list items
 -n X      When listing, only display X lines of description (default: 10)
 -g        When listing, use GLSA style (<ul>, <li>)
 -i regex  Use regex to select issues (default: "''' + DEFAULT_ISSUE_REGEX + '''")
 -t regex  Use regex to select TODOs  (default: "''' + DEFAULT_TODO_REGEX + '''")
 -a        Append line instead of replacing (applies to NOT-FOR-US, BUG)
 -T        Same as -t '^\s+TODO: check' (note the missing $)
 -s        Resort CVE list
 -u email  Username for Bugzilla (PyBugz) interface
 -p pass   Password for Bugzilla (PyBugz) interface
           (please note your PyBugz data is saved locally)

For each item, it guesses the name and searches the Portage Tree using `eix'.
After each action, it will print the entry as it would be saved.
'''
    miniusage()
    print "Please run ./bin/update prior to this script."



def main():
    import getopt
    try:
        optlist, list = getopt.getopt(sys.argv[1:],
                                  'ln:hi:t:Tau:p:sg')
    except getopt.GetoptError:
        usage(sys.argv[0])
        sys.exit(2)

    issue_regex = DEFAULT_ISSUE_REGEX
    todo_regex = DEFAULT_TODO_REGEX
    replace_line = True
    bugz_password = None
    bugz_username = None
    list_only = False
    glsa_style = False
    sort_only = False
    list_lines = 10
    for opt, arg in optlist:
        if opt == '-h':
            usage(sys.argv[0])
            sys.exit(0)
        if opt == '-n':
            list_lines = arg
        if opt == '-l':
            list_only = True
        if opt == '-i':
            issue_regex = arg.replace("{","(").replace("}",")").replace(",","|").replace(" ","")
        if opt == '-t':
            todo_regex = arg
        if opt == '-T':
            # no $ at the end
            todo_regex = '^\s+TODO: check'
        if opt == '-s':
            sort_only = True
        if opt == '-a':
            replace_line = False
        if opt == '-g':
            glsa_style = True
        if opt == '-p':
            bugz_password = arg
        if opt == '-u':
            bugz_username = arg

    if list_only and todo_regex == DEFAULT_TODO_REGEX and issue_regex != DEFAULT_ISSUE_REGEX:
        # reasoning: people enter "-l -i bla", which should just output that issue list
        # without further narrowing
        todo_regex = "."

    bugreporter = BugReporter(bugz_username, bugz_password)
    EntryEditor(issue_regex, todo_regex, replace_line, list_only, list_lines, glsa_style, sort_only, bugreporter)


if __name__ == "__main__":
    try:
        os.chdir(setup_paths())
        from cvetools import BugReporter, CVEData
        main()
    except KeyboardInterrupt:
        print '\n ! Exiting.'