victims-cve-db-cli.py

#!/usr/bin/python3

from argparse import ArgumentParser, HelpFormatter
from distutils.version import LooseVersion
import json
from os import path, walk
from os.path import join
import sys
from yaml import YAMLError, safe_load


class VictimsDB:
    def __init__(self, db_path=None):
        self._vulnerabilities = []
        self._db_path = db_path or path.join(path.dirname(path.abspath(__file__)),
                                             'database')

    @property
    def vulnerabilities(self):
        if not self._vulnerabilities:
            self.read_db()
        return self._vulnerabilities

    def read_db(self):
        for root, _, files in walk(self._db_path):
            for file in files:
                if not file.endswith(('.yaml', '.yml')):
                    continue
                with open(join(root, file), 'r') as stream:
                    try:
                        vulnerability = safe_load(stream)
                        self._vulnerabilities.append(vulnerability)
                    except YAMLError as exc:
                        print(exc)

    @staticmethod
    def is_version_affected(affected_versions, checked_version):
        checked_version = LooseVersion(checked_version)
        for version_range in affected_versions:
            operator = version_range[:2]
            if operator not in ('==', '<='):
                continue
            # https://github.com/victims/victims-cve-db#version-string-common
            if ',' in version_range:
                version = LooseVersion(version_range[2:].split(',')[0])
                series = LooseVersion(version_range.split(',')[1])
            else:
                version = LooseVersion(version_range[2:])
                series = None

            if operator == '==':
                if checked_version == version:
                    return True
            elif operator == '<=':
                if series:
                    if series <= checked_version <= version:
                        return True
                else:
                    if checked_version <= version:
                        return True
        return False

    def vulnerabilities_for_java_artifact(self, group_id, artifact_id, version):
        vulnerabilities = []
        for vulnerability in self.vulnerabilities:
            for affected in vulnerability.get('affected', []):
                if affected.get('groupId') == group_id and \
                   affected.get('artifactId') == artifact_id and \
                   self.is_version_affected(affected.get('version', []), version):
                    vulnerabilities.append(vulnerability)
        return vulnerabilities

    def vulnerabilities_for_python_artifact(self, name, version):
        vulnerabilities = []
        for vulnerability in self.vulnerabilities:
            for affected in vulnerability.get('affected', []):
                if affected.get('name') == name and \
                   self.is_version_affected(affected.get('version', []), version):
                    vulnerabilities.append(vulnerability)
        return vulnerabilities


class CLI(object):
    PROG = sys.argv[0]

    def __init__(self):
        self.parser = ArgumentParser(prog=self.PROG,
                                     description="Victims CVE Database CLI.",
                                     formatter_class=HelpFormatter)
        self.parser.add_argument("-d", "--db-path",
                                 help='alternate path to base directory with yaml files')
        self.parser.add_argument("-i", "--indent", action="store_true",
                                 help='pretty-print output json')

        subparsers = self.parser.add_subparsers()

        self.search_parser = subparsers.add_parser(
            'search',
            usage="%s [OPTIONS] search ..." % self.PROG,
            description="""
            Search Victims CVE Database for vulnerabilities in specified component.
            Example:
            ./victims-cve-db-cli.py -i search -e java -n org.apache.camel:camel-core -v 2.18.1"""
        )
        self.search_parser.add_argument("-e", "--ecosystem", choices=["java", "python"],
                                        help="java or python ecosystem")
        self.search_parser.add_argument("-n", "--name",
                                        help="name of component (groupId:artifactId for java)")
        self.search_parser.add_argument("-v", "--version",
                                        help="component version")

        self.search_parser.set_defaults(func=self.search)

    def search(self, args):
        if not args.ecosystem or not args.name or not args.version:
            self.parser.print_help()
            sys.exit(1)

        db = VictimsDB(args.db_path)
        if args.ecosystem == "java":
            try:
                group_id, artifact_id = args.name.split(':')
            except ValueError:
                print("Wrong groupId:artifactId specified")
                sys.exit(1)
            vulnerabilities = db.vulnerabilities_for_java_artifact(group_id=group_id,
                                                                   artifact_id=artifact_id,
                                                                   version=args.version)
            json.dump(vulnerabilities, sys.stdout, indent=args.indent)
        elif args.ecosystem == "python":
            vulnerabilities = db.vulnerabilities_for_python_artifact(name=args.name,
                                                                     version=args.version)
            json.dump(vulnerabilities, sys.stdout, indent=args.indent)

    def run(self):
        args = self.parser.parse_args()
        args.indent = 2 if args.indent else None
        if args.db_path and not path.isdir(args.db_path):
            print("%s is not a valid directory" % args.db_path)
            sys.exit(1)

        try:
            args.func(args)
        except AttributeError:
            if hasattr(args, 'func'):
                raise
            else:
                self.parser.print_help()


if __name__ == '__main__':
    cli = CLI()
    sys.exit(cli.run())