convert vulnerability YAML data (Grafeas format) to TSV

process.py

import csv
import sys
from operator import itemgetter
from types import SimpleNamespace

from yamlpath import Processor
from yamlpath import YAMLPath
from yamlpath.common import Parsers
from yamlpath.exceptions import YAMLPathException
from yamlpath.wrappers import ConsolePrinter
from yamlpath.wrappers import NodeCoords

loggingArgs = SimpleNamespace(quiet=True, verbose=True, debug=True)
logger = ConsolePrinter(loggingArgs)

# assign levels to known severity codes for sorting
SeverityCodes = ['', 'LOW', 'MEDIUM', 'HIGH', 'CRITICAL']
SeverityLevels = {code: i for (i, code) in enumerate(SeverityCodes)}

Columns = ['CVE ID', 'Effective Severity', 'Severity', 'Impacted Image',
           'Vulnerable Package', 'Remediated Package', 'URL']


def loadProcessor() -> Processor:
    yamlParser = Parsers.get_yaml_editor()

    # TODO: add support for stdin and stdout
    # TODO: try `sys.stdin`
    yamlFile = 'vulnerabilities.yaml'
    (yamlData, documentLoaded) = Parsers.get_yaml_data(yamlParser, logger,
                                                       yamlFile)
    if not documentLoaded:
        # error message already printed by yamlpath
        exit(1)

    processor = Processor(logger, yamlData)
    return processor


def main():
    rows = []

    processor = loadProcessor()

    try:
        impactedImage = NodeCoords.unwrap_node_coords(list(
            processor.get_nodes(YAMLPath('image_summary.tag'),
                                mustexist=True))[0])

        for node in processor.get_nodes(YAMLPath(
                'package_vulnerability_summary.vulnerabilities.*.*')):
            vulnerabilityData = NodeCoords.unwrap_node_coords(node)
            vulnerability = vulnerabilityData['vulnerability']

            # use `get()` for severity codes, because they may be missing
            severityCode = vulnerability.get('severity', '').upper()
            severityLevel = SeverityLevels.get(severityCode, 0)
            effectiveSeverityCode = vulnerability.get('effectiveSeverity',
                                                      '').upper()
            effectiveSeverityLevel = SeverityLevels.get(effectiveSeverityCode,
                                                        0)

            cveId = vulnerability['shortDescription']
            url = vulnerability['relatedUrls'][0]['url']

            packageIssue = vulnerability['packageIssue'][0]
            vulnerablePackage = f'{packageIssue["affectedPackage"]}/' \
                                f'{packageIssue["affectedVersion"]["fullName"]}'
            remediatedPackage = f'{packageIssue["fixedPackage"]}/' \
                                f'{packageIssue["fixedVersion"]["kind"]}'

            rows.append(
                [effectiveSeverityLevel, severityLevel, cveId,
                 effectiveSeverityCode, severityCode, impactedImage,
                 vulnerablePackage, remediatedPackage, url])

    except YAMLPathException as ex:
        logger.error(ex)

    rows.sort(key=itemgetter(0, 1),
              reverse=True)  # itemgetter is faster than lambda

    tsvWriter = csv.writer(sys.stdout, delimiter='\t')
    tsvWriter.writerow(Columns)
    for row in rows:
        tsvWriter.writerow(row[2:])  # skip the severity code sorting keys


if '__main__' == __name__:
    main()