#!/usr/bin/env python3

# Author: Jamie Strandboge <jamie@ubuntu.com>
# Copyright (C) 2005-2019 Canonical Ltd.
#
# This script is distributed under the terms and conditions of the GNU General
# Public License, Version 2 or later. See http://www.gnu.org/copyleft/gpl.html
# for details.

import datetime
import os
import re
import sys
import optparse
import cve_lib
import time
import json

import source_map

NANOSECONDS_PER_SECOND = 1000000000
SECONDS_PER_YEAR = 31556926
SECONDS_PER_DAY = 86400
POPULARITY_WINDOW = SECONDS_PER_YEAR * 1
NOW = time.time();

def which_source(pkg):
    for r in releases:
        if pkg in srcmap[r]:
            try:
                source = srcmap[r][pkg]['source']
                return source
            except KeyError:
                # package was in list for the release but there was no source
                # so most likely the source name is the same as the package name
                return pkg
    # this package wasn't in any source map most likely because it is 
    # packaged only for an EOL'd release such as precise
    if opt.debug:
        print("Not present in active release but being downloaded: ", pkg)
    return 'unknown'

def load_popularity_boost(popfile):
    with open(popfile) as json_data:
        d = json.load(json_data)
        length = len(d["results"][0]["series"][0]["values"])
        for x in range(0, length):
            timestamp = d["results"][0]["series"][0]["values"][x][0] / NANOSECONDS_PER_SECOND
            source = which_source(d["results"][0]["series"][0]["values"][x][1])
            downloads = d["results"][0]["series"][0]["values"][x][2]

            # Only look at universe packages and data in the last 1 years
            if (NOW - timestamp) > POPULARITY_WINDOW:
                continue;

            if source in popularity:
                    popularity[source] += downloads
            else:
                popularity[source] = downloads

    # Average downloads per package per day
    for source in popularity.keys():
        popularity[source] = popularity[source] / (POPULARITY_WINDOW /  SECONDS_PER_DAY)

    if opt.debug:
        print(popularity)

def invert_table(table):
    pkgs_to_cves = dict()
    for cve in table.keys():
        for pkg in table[cve]:
            pkgs_to_cves.setdefault(pkg, [])
            pkgs_to_cves[pkg].append(cve)

    return pkgs_to_cves

priorities = dict()
for p in ['untriaged'] + cve_lib.priorities:
    priorities.setdefault(p, dict())

packages = dict()
assignees = dict()
pkg_assignees = dict()
embargoed_packages = dict()
deferred_packages = dict()
popularity = dict()
pockets = { 'supported' : dict(),
            'universe' : dict(),
            'partner' : dict(),
            'only_in_supported_overlay' : dict() }

points = { 'negligible' : 0,
           'untriaged' : 5,
           'low' : 10,
           'medium' : 50,
           'high' : 100,
           'critical' : 200 }

# medium and higher get a point bump if older than too_old
too_old_packages = dict()
too_old = 30		# days
too_old_bump = 200	# bump for medium, high is this * 2, critical this * 4, low this / 20
today = time.time()

info = dict()

parser = optparse.OptionParser()
parser.add_option("-S", "--skip-devel", help="Show only those CVEs *not* in the current devel release", action="store_true")
parser.add_option("-D", "--only-devel", help="Show only those CVEs in the current devel release", action="store_true")
parser.add_option("-m", "--only-supported", help="Show only those CVEs that are supported", action="store_true")
parser.add_option("-u", "--not-supported", help="Show only those CVEs that aren't supported", action="store_true")
parser.add_option("-t", "--totals", help="Display totals", action="store_true")
parser.add_option("-p", "--packages", help="Report only on the given packages", action="append", type="string")
parser.add_option("-P", "--pkgfamily", help="Perform pkg family renamings (use argument multiple times for 'linux' and/or 'xen')", action="append", default=[])
parser.add_option("-X", "--exclude", help="Ignore specified packages", action="append", type="string")
parser.add_option("-r", "--release", help="Report only for the given releases", action="append", type="string")
parser.add_option("-d", "--debug", help="Report debug information while loading", action="store_true")
parser.add_option("-a", "--show-assigned", help="Show CVE assignments", action="store_true")
parser.add_option("--by-assignee", help="Show CVE assignments organized by assignee", action="store_true")
parser.add_option("--assignee", help="Show CVE assignments for assignee", action="append", type="string")
parser.add_option("--deferred-days", help="Don't show deferred items within NUM days", metavar="NUM", action="store", type="string")
parser.add_option("--skip-pending-overlay", help="Don't show items from the overlay ppa in the pending state", action="store_true")
parser.add_option("--skip-overlays", help="Don't show items from overlay ppas", action="store_true")
parser.add_option("--skip-snaps", help="Don't show snap items", action="store_true")
parser.add_option("--only-pending", help="Only show items in the pending state", action="store_true")
parser.add_option("--skip-pending", help="Don't show items in the pending state", action="store_true")
parser.add_option("-b", "--boost-popularity", help="Take package popularity into account when calculating priority score", action="store_true")
parser.add_option("--pbf", help="Popularity boost file", metavar="FILE")
parser.add_option("-j", "--json", help="Show more verbose output in JSON format", action="store_true")
(opt, args) = parser.parse_args()

map = source_map.load()
releases = cve_lib.all_releases
for eol in cve_lib.eol_releases:
    if eol in releases:
        releases.remove(eol)

if opt.skip_devel and cve_lib.devel_release != '':
    releases.remove(cve_lib.devel_release)

if opt.only_devel:
    releases = [cve_lib.devel_release]

if opt.release:
    releases = opt.release

if opt.boost_popularity:
    if opt.pbf:
        popularity_file = opt.pbf
    else:
        popularity_file = "package-popularity.json"
    srcmap = source_map.load('packages')
    load_popularity_boost(popularity_file)

(cves, uems) = cve_lib.get_cve_list()
(table, priority, cves, namemap, cveinfo) = cve_lib.load_table(cves, uems, opt)
date_pat = re.compile(r'^20[0-9][0-9]-[01][0-9]-[0-3][0-9]$')
if opt.deferred_days != None:
    deferred_today = datetime.datetime(int(datetime.date.today().year),
                                       int(datetime.date.today().month),
                                       int(datetime.date.today().day))
    deferred_delta = datetime.timedelta(days=int(opt.deferred_days))

for cve in sorted(cves):
    if not cve in table:
        continue

    # Load CVE if it isn't already cached
    if not cve in info:
        info.setdefault(cve, cve_lib.load_cve(cve_lib.find_cve(cve)))

    for pkg in sorted(table[cve].keys()):
        if opt.exclude and pkg in opt.exclude:
            continue

        supported = False
        partner = False
        universe = False
        found = False
        overlay = False
        archive = False
        for r in releases:
            if opt.skip_overlays and '/' in r:
                continue

            if opt.skip_snaps and r == 'snap':
                continue

            if r in table[cve][pkg]:
                if table[cve][pkg][r] in ['needed','deferred','pending','needs-triage','active']:
                    if cve_lib.is_supported(map, namemap[pkg][r], r, cveinfo[cve]):
                        if '/' in r:
                            overlay = True
                        else:
                            archive = True

                    # print >>sys.stderr, "%s for %s in %s (%s)" % (cve, pkg, r, table[cve][pkg][r])
                    if opt.only_pending and 'pending' not in table[cve][pkg][r]:
                        # print >>sys.stderr, "Skipping non-pending %s for %s in %s (%s)" % (cve, pkg, r, table[cve][pkg][r])
                        continue
                    elif opt.skip_pending_overlay and '/' in r and 'pending' in table[cve][pkg][r]:
                        # print >>sys.stderr, "Skipping pending %s for %s" % (cve, pkg)
                        continue
                    elif opt.skip_pending and 'pending' in table[cve][pkg][r]:
                        # print >>sys.stderr, "Skipping pending %s for %s" % (cve, pkg)
                        continue
                    elif opt.deferred_days != None and table[cve][pkg][r] == 'deferred':
                        if ("%s_comment" % r) in table[cve][pkg] and date_pat.search(table[cve][pkg]["%s_comment" % r]):
                            t = table[cve][pkg]["%s_comment" % r].split('-')
                            deferred_date = datetime.datetime(int(t[0]), int(t[1]), int(t[2]))
                            if deferred_today - deferred_delta > deferred_date:
                                deferred_packages[pkg] = 1
                            else:
                                continue
                        else:
                            continue

                    found = True
                    if cve_lib.is_supported(map, namemap[pkg][r], r, cveinfo[cve]):
                        supported = True
                    if cve_lib.is_partner(map, namemap[pkg][r], r):
                        partner = True
                    if cve_lib.is_universe(map, namemap[pkg][r], r, cveinfo[cve]):
                        universe = True

        if not found:
            continue

        if opt.only_supported and not supported and not partner:
            continue

        if opt.not_supported and (supported or partner):
            continue

        if supported:
            if pkg in pockets['supported']:
                pockets['supported'][pkg] += 1
            else:
                pockets['supported'][pkg] = 1

            if overlay and not archive:
                if pkg in pockets['only_in_supported_overlay']:
                    pockets['only_in_supported_overlay'][pkg] += 1
                else:
                    pockets['only_in_supported_overlay'][pkg] = 1

        if partner:
            if pkg in pockets['partner']:
                pockets['partner'][pkg] += 1
            else:
                pockets['partner'][pkg] = 1
        if universe:
            if pkg in pockets['universe']:
                pockets['universe'][pkg] += 1
            else:
                pockets['universe'][pkg] = 1

        if pkg in packages:
            packages[pkg] += 1
        else:
            packages[pkg] = 1

        p = priority[cve]['default']
        if pkg in priority[cve]:
            p = priority[cve][pkg]

        if pkg in priorities[p]:
            priorities[p][pkg] += 1
        else:
            priorities[p][pkg] = 1

        if os.path.islink('embargoed') and os.path.exists(os.path.join("embargoed", cve)):
            embargoed_packages[pkg] = "Undetermined"
            # CRD trumps PublicDate
            for tag in ['PublicDate', 'CRD']:
                crd = info[cve].get(tag, '')
                if crd != '':
                    embargoed_packages[pkg] = crd

        if p in ['low', 'medium', 'high', 'critical']:
            try:
                if cve_lib.cve_age(cve, info[cve]['PublicDate'], today) > too_old:
                    if pkg in too_old_packages:
                        if p in too_old_packages[pkg]:
                            too_old_packages[pkg][p] += 1
                        else:
                            too_old_packages[pkg][p] = 1
                    else:
                        too_old_packages[pkg] = dict()
                        too_old_packages[pkg][p] = 1
            except ValueError:
                continue

        if 'Assigned-to' in cveinfo[cve] and cveinfo[cve]['Assigned-to'] != "":
            if opt.show_assigned:
                if not pkg in pkg_assignees:
                    pkg_assignees[pkg] = []
                if cveinfo[cve]['Assigned-to'] not in pkg_assignees[pkg]:
                    pkg_assignees[pkg].append(cveinfo[cve]['Assigned-to'])
            if opt.by_assignee:
                if not (cveinfo[cve]['Assigned-to']) in assignees:
                    assignees[cveinfo[cve]['Assigned-to']] = dict()
                if not pkg in assignees[cveinfo[cve]['Assigned-to']]:
                    assignees[cveinfo[cve]['Assigned-to']][pkg] = []
                assignees[cveinfo[cve]['Assigned-to']][pkg].append(cve)


if opt.by_assignee:
    people = assignees.keys()
    people.sort()
    for p in people:
        if opt.assignee and p not in opt.assignee:
            continue
        sys.stdout.write("%s:\n" % p)
        pkgs = assignees[p].keys()
        pkgs.sort()
        for pkg in pkgs:
            sys.stdout.write("  %s:\n" % pkg)
            cves_p = assignees[p][pkg]
            cves_p.sort()
            for c in cves_p:
                sys.stdout.write("    %s\n" % c)
        sys.stdout.write("\n")
    sys.exit(0)

score_map = dict()
for pkg in sorted(packages.keys()):
    if opt.show_assigned and opt.assignee:
        if not pkg in pkg_assignees:
            continue
        found = False
        for p in opt.assignee:
            if p in pkg_assignees[pkg]:
                found = True
        if not found:
            continue

    score = 0
    for p in ['untriaged'] + cve_lib.priorities:
        if pkg in priorities[p]:
            score += priorities[p][pkg] * points[p]

    if pkg in too_old_packages:
        bump = 1
        for p in ['low', 'medium', 'high', 'critical']:
            if not p in too_old_packages[pkg]:
                continue
            if p == 'high':
                bump += too_old_packages[pkg][p] * too_old_bump * 2
            elif p == 'critical':
                bump += too_old_packages[pkg][p] * too_old_bump * 4
            elif p == 'low':
                bump += too_old_packages[pkg][p] * too_old_bump / 20
            else:
                bump += too_old_packages[pkg][p] * too_old_bump
        score += bump

    if opt.boost_popularity:
        boost = 1
        if pkg in popularity and popularity[pkg] > 0:
            boost = popularity[pkg]
        score *= boost

    score_map.setdefault(pkg, int(score))

if opt.json:
    pkgs_to_cves = invert_table(table)
    output = []
    for pkg in sorted(packages.keys()):
        if not pkg in score_map:
            continue
        package_entry = dict()
        package_entry.setdefault("package", pkg)
        package_entry.setdefault("score", score_map[pkg])
        package_entry.setdefault("popularity", "")
        if pkg in popularity:
            package_entry["popularity"] = str(popularity[pkg])

        package_entry.setdefault("cves", [])
        for cve in pkgs_to_cves[pkg]:
            package_entry["cves"].append({"cve": cve, "priority": cveinfo[cve]["Priority"]})
        output.append(package_entry)

    print(json.dumps(output, indent=4, sort_keys=True))
else:
    sys.stdout.write("Weight\tPackage Counts\n")
    sys.stdout.write("---------------------------------------------------------------------\n")

    for pkg in sorted(packages.keys()):
        if not pkg in score_map:
            continue

        if opt.totals:
            sys.stdout.write("%s: %s"%(pkg, packages[pkg]))
        else:
            sys.stdout.write("%s\t%s: %s total" % (str(score_map[pkg]), pkg, str(packages[pkg])))
            for p in ['untriaged'] + cve_lib.priorities:
                if pkg in priorities[p]:
                    sys.stdout.write(", %s %s" % (str(priorities[p][pkg]), p))

        extra_info = []
        if pkg in pockets['partner']:
            extra_info.append('PARTNER')
        elif pkg in pockets['supported']:
            if pkg in pockets['only_in_supported_overlay']:
                extra_info.append('SUPPORTED-ONLY-OVERLAY')
            else:
                extra_info.append('SUPPORTED')
        if pkg in embargoed_packages:
            extra_info.append('EMBARGOED:%s' % embargoed_packages[pkg])
        if pkg in deferred_packages:
            extra_info.append('CHECKDEFERRED')

        if len(extra_info) > 0:
            sys.stdout.write(" (%s)" % ",".join(extra_info))

        if opt.show_assigned and pkg in pkg_assignees:
            pkg_assignees[pkg].sort()
            sys.stdout.write(" (%s)" % ",".join(pkg_assignees[pkg]))

        sys.stdout.write("\n")