Refactor code to prepare for other output formats

Author: mariuszskon

cve_bin_tool/OutputEngine.py

@@ -24,33 +24,24 @@ def output_cves(self, outfile, output=None):
         """ Output a list of CVEs
         format is modules[checker_name][version] = dict{id: severity}
         """
-        # if the output is csv we must open a file
-        if output == "csv":
-            outfile = open(f"{outfile}", "w")
-
-            writer = csv.writer(outfile)
-            writer.writerow(["Module Name", "Version", "CVE Number", "Severity"])
 
         writer = csv.writer(outfile)
+        writer.writerow(["Module Name", "Version", "CVE Number", "Severity"])
         for modulename, versions in self.modules.items():
             for version, cve_list in versions.items():
                 for cve_number, cve_severity in cve_list.items():
                     row = [modulename, version, cve_number, cve_severity]
                     writer.writerow(row)
 
-        # We must also close the file
-        if output == "csv":
-            outfile.close()
-
-    def output_csv(self):
+    def output_file(self, output="csv"):
 
-        """ Generate a CSV file for list of CVE """
+        """ Generate a file for list of CVE """
 
         # Check if we need to generate a filename
         if self.filename == None:
-            self.generate_filename("csv")
+            self.generate_filename(output)
         else:
-            self.filename = f"{self.filename}.csv"
+            self.filename = f"{self.filename}.{output}"
             # check if the filename already exists
             file_list = os.listdir(os.getcwd())
             if self.filename in file_list:
@@ -60,7 +51,7 @@ def output_csv(self):
                 self.logger.info(
                     "Generating a new filename with Default Naming Convention"
                 )
-                self.generate_filename("csv")
+                self.generate_filename(output)
 
             # try opening that file
             try:
@@ -70,10 +61,11 @@ def output_csv(self):
             except Exception as E:
                 self.logger.warning(E)
                 self.logger.info("Switching Back to Default Naming Convention")
-                self.generate_filename("csv")
+                self.generate_filename(output)
 
         # Log the filename generated
-        self.logger.info(f"CSV output stored at {os.getcwd()}/{self.filename}")
+        self.logger.info(f"Output stored at {os.getcwd()}/{self.filename}")
 
         # call to output_cves
-        self.output_cves(self.filename, "csv")
+        with open(self.filename, "w") as f:
+            self.output_cves(f, output)

cve_bin_tool/cli.py

@@ -491,10 +491,9 @@ def main(argv=None, outfile=sys.stdout):
                     and args.output == "console"
                 ):
                     output.output_cves(outfile)
-
-                # If the args are passed for csv we will generate a CSV output
-                if args.output == "csv":
-                    output.output_csv()
+                else:
+                    # If the args are passed for csv/json we will generate a file output
+                    output.output_file(args.output)
 
             # Use the number of files with known cves as error code
             # as requested by folk planning to automate use of this script.