#!/usr/bin/env python

import sys
import os
import re
import anchore.anchore_utils

def perform_cve_scan_simple(severity='all'):
    ret = {}
    warns = list()

    pkgmaps = {}
    imgmaps = {}
    flvmaps = {}

    cve_data_context = {}
    for imageId in config['images']:

        distrometa = anchore.anchore_utils.get_distro_from_imageId(imageId)
        idistro = distrometa['DISTRO']
        idistrovers = distrometa['DISTROVERS']
        distrodict = anchore.anchore_utils.get_distro_flavor(idistro, idistrovers)
        flavor = distrodict['flavor']

        try:
            dstr, new_cve_data = anchore.anchore_utils.cve_load_data_imageId(imageId, cve_data_context=cve_data_context)

            if dstr not in pkgmaps:
                pkgmaps[dstr] = {}

            if dstr not in imgmaps:
                imgmaps[dstr] = list()
                
            if dstr not in flvmaps:
                flvmaps[dstr] = flavor

            imgmaps[dstr].append(imageId)

            pkgmap = pkgmaps[dstr]

            pkgs = anchore.anchore_utils.load_analysis_output(imageId, 'package_list', 'pkgs.all')
            pkgs.update(anchore.anchore_utils.load_analysis_output(imageId, 'package_list', 'pkgs_plus_source.all'))

            for pkg in pkgs.keys():
                vers = pkgs[pkg]

                if pkg not in pkgmap:
                    pkgmap[pkg] = {'versions':{}}

                if vers not in pkgmap[pkg]['versions']:
                    pkgmap[pkg]['versions'][vers] = list()

                pkgmap[pkg]['versions'][vers].append(imageId)

        except Exception as err:
            warns.append(config['meta']['shortId']+" ("+ config['meta']['humanname'] +") cannot scan image: no CVE data available for distro ("+str(idistro)+":"+str(idistrovers)+")")

    for dstr in pkgmaps.keys():
        cve_results = anchore.anchore_utils.cve_scanimages(imgmaps[dstr], pkgmaps[dstr], flvmaps[dstr], cve_data_context[dstr])

        if cve_results:
            for cveId in cve_results.keys():
                if severity == 'all' or cve_results[cveId]['severity'] == severity:
                    if cveId not in ret:
                        ret[cveId] = {'severity':'Unknown', 'imageIds':list()}

                    outsev = "N/A"
                    if 'severity' in cve_results[cveId]:
                        outsev = cve_results[cveId]['severity']

                    outurl = "N/A"
                    if 'url' in cve_results[cveId]:
                        outurl = cve_results[cveId]['url']

                    ret[cveId]['severity'] = outsev
                    ret[cveId]['url'] = outurl
                    ret[cveId]['imageIds'] = ret[cveId]['imageIds'] + cve_results[cveId]['images']
        
    return(ret, warns)

# 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(["*CVE_ID","Severity","Affected_Count", "Affected_Image_Ids", "URL"])

# perform query
result = {}
warns = list()

try:
    result, warns = perform_cve_scan_simple('all')
except Exception as err:
    warns.append("could not perform CVE scan due to unexpected error: " + str(err))

for k in result.keys():
    if 'all' in config['params'] or result[k]['severity'] in config['params']:
        outlist.append([k, result[k]['severity'], str(len(result[k]['imageIds'])), ' '.join(result[k]['imageIds']), result[k]['url']])

# save output list as output file for anchore to ingress
anchore.anchore_utils.write_kvfile_fromlist(config['output'], outlist)
if len(warns) > 0:
    anchore.anchore_utils.write_plainfile_fromlist(config['output_warns'], warns)

# cleanup and exit
sys.exit(0)
