Load premium database file if available

SCH227 · Jan 9, 2021 · a96ed71 · a96ed71
1 parent b289752
commit a96ed71
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/safety/cli.py b/safety/cli.py
@@ -31,6 +31,8 @@ def cli():
 @click.option("--full-report/--short-report", default=False,
               help='Full reports include a security advisory (if available). Default: '
                    '--short-report')
+# @click.option("--cvss", default="",
+#               help="Extend reports with premium CVSS information (if available). Default: empty")
 @click.option("--bare/--not-bare", default=False,
               help='Output vulnerable packages only. '
                    'Useful in combination with other tools. '

diff --git a/safety/safety.py b/safety/safety.py
@@ -107,14 +107,15 @@ def fetch_database_file(path, db_name):
         return json.loads(f.read())
 
 
-def fetch_database(full=False, key=False, db=False, cached=False, proxy={}):
+def fetch_database(full=False, premium=False, key=False, db=False, cached=False, proxy={}):
 
     if db:
         mirrors = [db]
     else:
         mirrors = API_MIRRORS if key else OPEN_MIRRORS
 
     db_name = "insecure_full.json" if full else "insecure.json"
+    db_name = "insecure_premium.json" if premium else db_name
     for mirror in mirrors:
         # mirror can either be a local path or a URL
         if mirror.startswith("http://") or mirror.startswith("https://"):
@@ -137,6 +138,7 @@ def check(packages, key, db_mirror, cached, ignore_ids, proxy):
     key = key if key else os.environ.get("SAFETY_API_KEY", False)
     db = fetch_database(key=key, db=db_mirror, cached=cached, proxy=proxy)
     db_full = None
+    premium = False
     vulnerable_packages = frozenset(db.keys())
     vulnerable = []
     for pkg in packages:
@@ -154,9 +156,14 @@ def check(packages, key, db_mirror, cached, ignore_ids, proxy):
                 spec_set = SpecifierSet(specifiers=specifier)
                 if spec_set.contains(pkg.version):
                     if not db_full:
-                        db_full = fetch_database(full=True, key=key, db=db_mirror, cached=cached, proxy=proxy)
+                        premium = True
+                        db_full = fetch_database(full=True, premium=True, key=key, db=db_mirror, cached=cached, proxy=proxy)
+                        if not db_full:
+                            premium = False
+                            db_full = fetch_database(full=True, key=key, db=db_mirror, cached=cached, proxy=proxy)
                     for data in get_vulnerabilities(pkg=name, spec=specifier, db=db_full):
                         vuln_id = data.get("id").replace("pyup.io-", "")
+                        cve_ids = map(str.strip, data.get("cve").split(','))
                         if vuln_id and vuln_id not in ignore_ids:
                             vulnerable.append(
                                 Vulnerability(