#!/usr/bin/env python

import sys
import os
import re
import anchore.anchore_utils

def perform_cve_scan(cvedatadir='/tmp', severity='all'):
    severity = severity
    ret = {}
    cve_ret = {}

    for imageId in config['images']:
        image = allimages[imageId]
        cve_data = anchore.anchore_utils.cve_load_data(cvedatadir, image)

        earliest_vulnerable = None

        if cve_data:
            cve_results = anchore.anchore_utils.cve_scanimage(cve_data, image)

            if cve_results:
                for cveId in cve_results.keys():
                    if severity == 'all' or cve_results[cveId]['severity'] == severity:
                        depth=0
                        branch = list()
                        for fid in image.get_familytree():
                            fimage = allimages[fid]
                            if fimage.is_user() or fimage.is_base() or fid == image.meta['imageId']:
                                if fid == image.meta['imageId']:
                                    fcve_results = cve_results.copy()
                                else:
                                    fcve_results = anchore.anchore_utils.cve_scanimage(cve_data, fimage)

                                if cveId in fcve_results.keys():
                                    branch.append(fid)
                                    cve_ret = cve_ret.copy()
                                    cve_ret.update(fcve_results)

                        if len(branch) > 0:
                            if cveId not in ret:
                                ret[cveId] = list()
                            ret[cveId].append(branch)

    branches = ret
    cve_results = cve_ret

    ret = {}
    for cveId in branches.keys():

        tofix = {}
        torebuild = {}

        for b in branches[cveId]:
            image = allimages[b[0]]

            tofix[''.join(image.get_human_string().split())] = True
            for bk in b[1:]:
                image = allimages[bk]
                torebuild[''.join(image.get_human_string().split())] = True

            total_vuln = len(tofix.keys()) + len(torebuild.keys())

            if not cveId in ret:
                ret[cveId] = list()

            if len(torebuild.keys()) <= 0:
                rebuildstr = 'None'
            else:
                rebuildstr = ','.join(torebuild.keys())

            if len(tofix.keys()) <= 0:
                fixstr = 'None'
            else:
                fixstr = ','.join(tofix.keys())

        if cve_results[cveId]['fixVers'] == 'None':
            fixpkg = "None"
        else:
            fixpkg = cve_results[cveId]['pkgName'] + "-" + cve_results[cveId]['fixVers']
        ret[cveId].append([cveId, cve_results[cveId]['severity'], str(total_vuln), cve_results[cveId]['pkgName'] + "-" + cve_results[cveId]['imageVers'], fixpkg, fixstr, rebuildstr ])
    return(ret)

# main routine

# parse cmdline into config dict
try:
    config = anchore.anchore_utils.init_query_cmdline(sys.argv, "params: <Low|Medium|High|Critical|Unknown|Negligible> <another level> ...\nhelp: use 'all' to scan for all severity levels")
except Exception as err:
    print str(err)
    sys.exit(1)

# this means that no config is available, for example when 'help' is given as a parameter
if not config:
    sys.exit(0)

# handle default case if no params are given
if len(config['params']) <= 0:
    config['params'] = ['all']

# set up output list
outlist = list()
outlist.append(["CVEID","Severity","*TotalAffected","VulnerablePackage", "FixAvailable","FixImages","RebuildImages"])

# perform query
allimages = {}
for imageId in config['images']:
    image = anchore.anchore_image.AnchoreImage(imageId, config['anchore_config']['image_data_store'], allimages)

    if imageId not in allimages:
        allimages[imageId] = image
        
for sev in config['params']:
    try:
        cvedatadir = '/'.join([config['anchore_config']['image_data_store'], '..', 'cve-data'])
        result = perform_cve_scan(cvedatadir, sev)

        for k in result.keys():
            for kk in result[k]:
                outlist.append(kk)
    except:
        pass

# save output list as output file for anchore to ingress
anchore.anchore_utils.write_kvfile_fromlist(config['output'], outlist)

# cleanup and exit
allimages.clear()
sys.exit(0)
