#!/usr/bin/env python

import sys
import os
import re
import anchore.anchore_utils

def perform_cve_scan(severity='all'):
    severity = severity
    ret = {}
    cve_ret = {}
    warns = list()
    
    cve_data_context = {}
    
    for imageId in config['images']:
        image = allimages[imageId]
        try:
            cve_data = anchore.anchore_utils.cve_load_data(image, cve_data_context=cve_data_context)
        except ValueError as err:
            warns.append(image.meta['shortId']+" ("+image.get_human_name()+") " + str(err))
            cve_data = None

        earliest_vulnerable = None
        if cve_data:
            cve_results = anchore.anchore_utils.cve_scanimage(cve_data, imageId)

            fresults = {}
            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)
                    fresults[fid] = fcve_results                
            print "finished scanning"
            if cve_results:
                for cveId in cve_results.keys():
                    for cvepkg in cve_results[cveId]:
                        if severity == 'all' or cvepkg['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']:
                                    fcve_results = fresults[fid]
                                    if cveId in fcve_results.keys():
                                        branch.append(fid)
                                        if not cve_ret:
                                            cve_ret = fcve_results.copy()
                                        cve_ret.update(fcve_results)

                            if len(branch) > 0:
                                if cveId not in ret:
                                    ret[cveId] = list()
                                ret[cveId].append(branch)
        else:
            warns.append(image.meta['shortId']+" ("+image.get_human_name()+") cannot scan image: no CVE data loaded from feed")
                             
    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())
                rebuildstr = ' '.join(torebuild.keys())

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

        if cveId in cve_results:
            for cvepkg in cve_results[cveId]:
                if cvepkg['fixVers'] == 'None':
                    fixpkg = "None"
                else:
                    fixpkg = cvepkg['fixVers']

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

                outpkg = "N/A"
                if 'pkgName' in cvepkg and 'imageVers' in cvepkg:
                    outpkg = cvepkg['pkgName'] + "-" + cvepkg['imageVers']

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

                ret[cveId].append([str(cveId), str(outsev), str(total_vuln), str(outpkg), str(fixpkg), str(fixstr), str(rebuildstr), str(outurl)])

    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","*Total_Affected","Vulnerable_Package", "Fix_Available","Fix_Images","Rebuild_Images", "URL"])

# 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'])
    warns = list()
    result = {}
    try:
        #result, warns = perform_cve_scan(cvedatadir, sev)
        result, warns = perform_cve_scan('all')
    except ValueError as err:
        print str(err)
        #sys.exit(1)

    for k in result.keys():
        for kk in result[k]:
            if 'all' in config['params'] or kk[1] in config['params']:
                outlist.append(kk)

except Exception as err:
    print str(err)
    import traceback
    traceback.print_exc()

# 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
allimages.clear()
sys.exit(0)
