several optimization in cvedb and csv2cve files and fixes issue #413 (#477)

Author: Niraj-Kamdar

cve_bin_tool/csv2cve.py

@@ -53,11 +53,12 @@ def csv2cve(filename):
             LOGGER.error("Error: invalid CSV")
             return ERR_BADCSV
 
-        required_columns = ["vendor", "product", "version"]
-        for column in required_columns:
-            if column not in csvdata.fieldnames:
-                LOGGER.error(f"Error: no {column} column found")
-                return ERR_MISSINGCOLUMN
+        required_columns = {"vendor", "product", "version"}
+        csv_columns = set(csvdata.fieldnames)
+        missing_columns = required_columns - csv_columns
+        if missing_columns != set():
+            LOGGER.error(f"Error: no {missing_columns} columns found")
+            return ERR_MISSINGCOLUMN
 
         # Initialize the NVD database
         cvedb = CVEDB()
@@ -70,8 +71,14 @@ def csv2cve(filename):
             )
             cves = cvedb.get_cves(row["vendor"], row["product"], row["version"])
             if cves:
-                LOGGER.info("\n".join(sorted(cves.keys())))
-                cveoutput.append(cves.keys())
+                s = map(
+                    lambda key: " ".join(
+                        [row["vendor"], row["product"], row["version"], key, cves[key]]
+                    ),
+                    sorted(cves.keys()),
+                )
+                LOGGER.info("\n".join(s))
+                cveoutput.append(cves)
             else:
                 LOGGER.debug("No CVEs found. Is the vendor/product info correct?")
             LOGGER.debug("")

cve_bin_tool/cvedb.py

@@ -23,7 +23,8 @@
     import urllib2 as request
 
 from collections import namedtuple
-from .log import LOGGER
+from string import ascii_lowercase
+from cve_bin_tool.log import LOGGER
 from pkg_resources import parse_version
 
 logging.basicConfig(level=logging.DEBUG)
@@ -268,10 +269,7 @@ def get_cves(self, vendor, product, version):
         query = """SELECT CVE_number FROM cve_range WHERE
         vendor=? AND product=? AND version=?"""
         cursor.execute(query, [vendor, product, version])
-        # FIXME: this seems inefficient
-        cve_list = []
-        for cve in cursor:
-            cve_list.append(cve[0])
+        cve_list = list(map(lambda x: x[0], cursor.fetchall()))
 
         # Check for any ranges
         query = """SELECT CVE_number, versionStartIncluding, versionStartExcluding, versionEndIncluding, versionEndExcluding FROM cve_range WHERE
@@ -297,6 +295,7 @@ def get_cves(self, vendor, product, version):
                 versionEndExcluding = self.openssl_convert(versionEndExcluding)
 
             parsed_version = parse_version(version)
+
             # check the start range
             passes_start = False
             if (
@@ -336,21 +335,17 @@ def get_cves(self, vendor, product, version):
             ):
                 # then there is no end range so it passes
                 passes_end = True
-
             # if it fits into both ends of the range, add the cve number
             if passes_start and passes_end:
                 cve_list.append(cve_number)
 
         # Go through and get all the severities
         if cve_list:
-            query = f'SELECT CVE_number, severity from cve_severity where CVE_number IN ({",".join("?" for i in cve_list)}) ORDER BY CVE_number ASC'
+            query = f'SELECT CVE_number, severity from cve_severity where CVE_number IN ({",".join(["?"]*len(cve_list))}) ORDER BY CVE_number ASC'
             cursor.execute(query, cve_list)
             # Everything expects a data structure of cve[number] = severity so you can search through keys
             # and do other easy manipulations
-            severity_list = dict()
-            for cve_id, severity in cursor:
-                severity_list[cve_id] = severity
-            return severity_list
+            return dict(cursor)
 
         return cve_list
 
@@ -361,10 +356,7 @@ def openssl_convert(self, version):
             return version
 
         lastchar = version[len(version) - 1]
-        letters = {
-            letter: str(index)
-            for index, letter in enumerate("abcdefghijklmnopqrstuvwxyz")
-        }
+        letters = dict(zip(ascii_lowercase, range(26)))
 
         if lastchar in letters:
             version = f"{version[0 : len(version) - 1]}.{letters[lastchar]}"