Unified csv & console output (#410)

* Added OutputEngine to group together cve output options
* Updated cli.py to use it, and to have flags for output type and output filename
* fixes #409

Author: SinghHrmn

cve_bin_tool/OutputEngine.py

@@ -0,0 +1,79 @@
+import csv
+import os
+
+from datetime import datetime
+from .log import LOGGER
+
+
+class OutputEngine(object):
+    def __init__(self, filename=None, modules=None, logger=None):
+        if logger is None:
+            logger = LOGGER.getChild(self.__class__.__name__)
+        self.filename = filename
+        self.modules = modules
+        self.logger = logger
+
+    def generate_filename(self, extention=None):
+        """ Generates a random filename"""
+        if not extention == None:
+            now = datetime.now().strftime("%Y-%m-%d.%H:%m:%S")
+            self.filename = f"output.cve-bin-tool.{now}.{extention}"
+
+    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)
+        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):
+
+        """ Generate a CSV file for list of CVE """
+
+        # Check if we need to generate a filename
+        if self.filename == None:
+            self.generate_filename("csv")
+        else:
+            self.filename = f"{self.filename}.csv"
+            # check if the filename already exists
+            file_list = os.listdir(os.getcwd())
+            if self.filename in file_list:
+                self.logger.warning(
+                    f"Failed to write at '{self.filename}'. File already exists"
+                )
+                self.logger.info(
+                    "Generating a new filename with Default Naming Convention"
+                )
+                self.generate_filename("csv")
+
+            # try opening that file
+            try:
+                with open(self.filename, "w") as f:
+                    f.write("testing")
+                os.remove(self.filename)
+            except Exception as E:
+                self.logger.warning(E)
+                self.logger.info("Switching Back to Default Naming Convention")
+                self.generate_filename("csv")
+
+        # Log the filename generated
+        self.logger.info(f"CSV output stored at {os.getcwd()}/{self.filename}")
+
+        # call to output_cves
+        self.output_cves(self.filename, "csv")

cve_bin_tool/cli.py

@@ -29,6 +29,7 @@
 from .extractor import Extractor
 from .strings import Strings
 from .file import is_binary
+from .OutputEngine import OutputEngine
 
 from .cvedb import CVEDB
 from .log import LOGGER
@@ -291,18 +292,6 @@ def scan_files_unpack(unpacked):
     return scan_files(*unpacked)
 
 
-def output_cves(outfile, modules):
-    """ Output a list of CVEs
-    format is modules[checker_name][version] = dict{id: severity}
-    """
-    writer = csv.writer(outfile)
-    for modulename, versions in 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)
-
-
 def main(argv=None, outfile=sys.stdout):
     """ Scan a binary file for certain open source libraries that may have CVEs """
     if argv is None:
@@ -338,6 +327,21 @@ def main(argv=None, outfile=sys.stdout):
         action=LogAction,
         choices=["debug", "info", "warning", "error", "critical"],
     )
+    output_group.add_argument(
+        "-o",
+        "--output",
+        action="store",
+        choices=["csv", "json", "console"],
+        default="console",
+        help="update output format (default: console)",
+    )
+    output_group.add_argument(
+        "-f",
+        "--filename",
+        action="store",
+        default=None,
+        help="provide output file name",
+    )
     parser.add_argument(
         "-v", "--version", action="version", version=f"{get_version_string()}",
     )
@@ -478,8 +482,19 @@ def main(argv=None, outfile=sys.stdout):
                     )
                 )
                 LOGGER.info(f"Known CVEs in {affected_string}:")
-                if LOGGER.getEffectiveLevel() != logging.CRITICAL:
-                    output_cves(outfile, scanner.all_cves)
+
+                # Creates a Object for OutputEngine
+                output = OutputEngine(modules=scanner.all_cves, filename=args.filename)
+
+                if (
+                    LOGGER.getEffectiveLevel() != logging.CRITICAL
+                    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()
 
             # Use the number of files with known cves as error code
             # as requested by folk planning to automate use of this script.