From 9eb06046c203a017d3cad83d57fdacff8c666261 Mon Sep 17 00:00:00 2001
From: Sam Washko <sam@protectai.com>
Date: Wed, 13 Mar 2024 09:58:10 -0700
Subject: [PATCH] Add generate_report() (#114)

Co-authored-by: Sean Morgan <seanmorgan91@gmail.com>
---
 modelscan/cli.py       | 27 ++++++++-------------------
 modelscan/modelscan.py | 26 ++++++++++++++++++++++++++
 modelscan/reports.py   |  4 ++--
 3 files changed, 36 insertions(+), 21 deletions(-)

diff --git a/modelscan/cli.py b/modelscan/cli.py
index 1005b8d..230b083 100644
--- a/modelscan/cli.py
+++ b/modelscan/cli.py
@@ -1,7 +1,6 @@
 import logging
 import sys
 import os
-import importlib
 from pathlib import Path
 from typing import Optional, Dict, Any
 from tomlkit import parse
@@ -9,7 +8,6 @@
 import click
 
 from modelscan.modelscan import ModelScan
-from modelscan.reports import Report
 from modelscan._version import __version__
 from modelscan.settings import (
     SettingsUtils,
@@ -133,25 +131,16 @@ def scan(
     else:
         raise click.UsageError("Command line must include a path")
 
-    report_settings: Dict[str, Any] = {}
-    if reporting_format == "custom":
-        reporting_module = settings["reporting"]["module"]  # type: ignore[index]
-    else:
-        reporting_module = DEFAULT_REPORTING_MODULES[reporting_format]
-
-    report_settings = settings["reporting"]["settings"]  # type: ignore[index]
-    report_settings["show_skipped"] = show_skipped
-    report_settings["output_file"] = output_file
+    # Report scan results
+    if reporting_format is not "custom":
+        modelscan._settings["reporting"]["module"] = DEFAULT_REPORTING_MODULES[
+            reporting_format
+        ]
 
-    try:
-        (modulename, classname) = reporting_module.rsplit(".", 1)
-        imported_module = importlib.import_module(name=modulename, package=classname)
+    modelscan._settings["reporting"]["settings"]["show_skipped"] = show_skipped
+    modelscan._settings["reporting"]["settings"]["output_file"] = output_file
 
-        report_class: Report = getattr(imported_module, classname)
-        report_class.generate(scan=modelscan, settings=report_settings)
-
-    except Exception as e:
-        logger.error(f"Error generating report using {reporting_module}: {e}")
+    modelscan.generate_report()
 
     # exit code 3 if no supported files were passed
     if not modelscan.scanned:
diff --git a/modelscan/modelscan.py b/modelscan/modelscan.py
index 7627fac..b24a5f9 100644
--- a/modelscan/modelscan.py
+++ b/modelscan/modelscan.py
@@ -289,6 +289,32 @@ def is_compatible(self, path: str) -> bool:
 
         return False
 
+    def generate_report(self) -> Optional[str]:
+        reporting_module = self._settings["reporting"]["module"]
+        report_settings = self._settings["reporting"]["settings"]
+
+        scan_report = None
+        try:
+            (modulename, classname) = reporting_module.rsplit(".", 1)
+            imported_module = importlib.import_module(
+                name=modulename, package=classname
+            )
+
+            report_class = getattr(imported_module, classname)
+            scan_report = report_class.generate(scan=self, settings=report_settings)
+
+        except Exception as e:
+            logger.error(f"Error generating report using {reporting_module}: {e}")
+            self._errors.append(
+                ModelScanError(
+                    "ModelScan",
+                    ErrorCategories.MODEL_SCAN,
+                    f"Error generating report using {reporting_module}: {e}",
+                )
+            )
+
+        return scan_report
+
     @property
     def issues(self) -> Issues:
         return self._issues
diff --git a/modelscan/reports.py b/modelscan/reports.py
index a6160ab..1fbf246 100644
--- a/modelscan/reports.py
+++ b/modelscan/reports.py
@@ -88,12 +88,12 @@ def generate(
         settings: Dict[str, Any] = {},
     ) -> None:
         report: Dict[str, Any] = scan._generate_results()
-        if not settings["show_skipped"]:
+        if not settings.get("show_skipped"):
             del report["summary"]["skipped"]
 
         print(json.dumps(report))
 
-        output = settings["output_file"]
+        output = settings.get("output_file")
         if output:
             with open(output, "w") as outfile:
                 json.dump(report, outfile)