From e420d4abaa75150e8ad9a7123a022096bfcdd4cd Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 12 Jan 2024 15:08:02 -0800 Subject: [PATCH 1/5] add json to reports.py --- modelscan/cli.py | 14 +++++++------ modelscan/reports.py | 47 ++++++++++++++++++++++++++++++------------- modelscan/settings.py | 8 +++++--- 3 files changed, 46 insertions(+), 23 deletions(-) diff --git a/modelscan/cli.py b/modelscan/cli.py index 7383d7d..b5bd5ff 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -111,12 +111,14 @@ def scan( modelscan.scan(pathlibPath) else: raise click.UsageError("Command line must include a path") - ConsoleReport.generate( - modelscan.issues, - modelscan.errors, - modelscan._skipped, - show_skipped=show_skipped, - ) + + # get reporting module + + # ConsoleReport.generate( + # scan=modelscan, + # show_skipped=show_skipped, + # settings=settings["reporting"] + # ) # exit code 3 if no supported files were passed if not modelscan.scanned: diff --git a/modelscan/reports.py b/modelscan/reports.py index f0e9fdb..721e05a 100644 --- a/modelscan/reports.py +++ b/modelscan/reports.py @@ -1,9 +1,11 @@ import abc import logging -from typing import List, Optional +import json +from typing import Optional, Dict, Any from rich import print +from modelscan.modelscan import ModelScan from modelscan.error import Error from modelscan.issues import Issues, IssueSeverity @@ -20,10 +22,9 @@ def __init__(self) -> None: @staticmethod def generate( - issues: Issues, - errors: List[Error], - skipped: List[str], + scan: ModelScan, show_skipped: bool = False, + settings: Dict[str, Any] = {}, ) -> Optional[str]: """ Generate report for the given codebase. @@ -39,14 +40,13 @@ def generate( class ConsoleReport(Report): @staticmethod def generate( - issues: Issues, - errors: List[Error], - skipped: List[str], + scan: ModelScan, show_skipped: bool = False, + settings: Dict[str, Any] = {}, ) -> None: - issues_by_severity = issues.group_by_severity() + issues_by_severity = scan.issues.group_by_severity() print("\n[blue]--- Summary ---") - total_issue_count = len(issues.all_issues) + total_issue_count = len(scan.issues.all_issues) if total_issue_count > 0: print(f"\nTotal Issues: {total_issue_count}") print(f"\nTotal Issues By Severity:\n") @@ -66,18 +66,37 @@ def generate( else: print("\n[green] No issues found! 🎉") - if len(errors) > 0: + if len(scan.errors) > 0: print("\n[red]--- Errors --- ") - for index, error in enumerate(errors): + for index, error in enumerate(scan.errors): print(f"\nError {index+1}:") print(str(error)) - if len(skipped) > 0: + if len(scan.skipped) > 0: print("\n[blue]--- Skipped --- ") print( - f"\nTotal skipped: {len(skipped)} - run with --show-skipped to see the full list." + f"\nTotal skipped: {len(scan.skipped)} - run with --show-skipped to see the full list." ) if show_skipped: print(f"\nSkipped files list:\n") - for file_name in skipped: + for file_name in scan.skipped: print(str(file_name)) + + +class JSONReport(Report): + @staticmethod + def generate( + scan: ModelScan, + show_skipped: bool = False, + settings: Dict[str, Any] = {}, + ) -> None: + report: Dict[str, Any] = scan._generate_results() + if not show_skipped: + del report["skipped"] + + print(json.dumps(report)) + + output = settings["output_file"] + if output: + with open(output, "w") as outfile: + json.dump(report, outfile) diff --git a/modelscan/settings.py b/modelscan/settings.py index a877046..8ced0a1 100644 --- a/modelscan/settings.py +++ b/modelscan/settings.py @@ -1,8 +1,5 @@ import tomlkit -from typing import Any - - DEFAULT_SCANNERS = [ "modelscan.scanners.H5Scan", "modelscan.scanners.KerasScan", @@ -11,6 +8,7 @@ "modelscan.scanners.PickleScan", "modelscan.scanners.PyTorchScan", ] +from typing import Any DEFAULT_SETTINGS = { "supported_zip_extensions": [".zip", ".npz"], @@ -95,6 +93,10 @@ "MEDIUM": {}, "LOW": {}, }, + "reporting_module": { + "module": "modelscan.reports.ConsoleReport", + "settings": {}, + }, }, } From 159cbcb7511fd1c0381d3c682551afbaa0bf6e81 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 12 Jan 2024 16:05:46 -0800 Subject: [PATCH 2/5] add CLI option for json or custom reporting from toml --- modelscan/cli.py | 51 +++++++++++++++++++++++++++++++++++-------- modelscan/reports.py | 7 ++---- modelscan/settings.py | 9 ++++---- 3 files changed, 49 insertions(+), 18 deletions(-) diff --git a/modelscan/cli.py b/modelscan/cli.py index b5bd5ff..1ffd394 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -1,16 +1,21 @@ import logging import sys import os +import importlib from pathlib import Path -from typing import Optional +from typing import Optional, Dict, Any from tomlkit import parse import click from modelscan.modelscan import ModelScan -from modelscan.reports import ConsoleReport +from modelscan.reports import Report from modelscan._version import __version__ -from modelscan.settings import SettingsUtils, DEFAULT_SETTINGS +from modelscan.settings import ( + SettingsUtils, + DEFAULT_SETTINGS, + AVAILABLE_REPORTING_MODULES, +) from modelscan.tools.cli_utils import DefaultGroup logger = logging.getLogger("modelscan") @@ -69,6 +74,20 @@ def cli() -> None: type=click.Path(exists=True, dir_okay=False), help="Specify a settings file to use for the scan. Defaults to ./modelscan-settings.toml.", ) +@click.option( + "-f", + "--format", + type=click.Choice(["console", "json", "custom"]), + default="console", + help="Format of the output. Options are console or json (default: console)", +) +@click.option( + "-o", + "--output-file", + type=click.Path(), + default=None, + help="Optional json reporting output file", +) @cli.command( help="[Default] Scan a model file or diretory for ability to execute suspicious actions. " ) # type: ignore @@ -79,6 +98,8 @@ def scan( path: Optional[str], show_skipped: bool, settings_file: Optional[str], + format: str, + output_file: Path, ) -> int: logger.setLevel(logging.INFO) logger.addHandler(logging.StreamHandler(stream=sys.stdout)) @@ -112,13 +133,25 @@ def scan( else: raise click.UsageError("Command line must include a path") - # get reporting module + report_settings: Dict[str, Any] = {} + if format == "custom": + reporting_module = settings["reporting"]["module"] # type: ignore[index] + report_settings = settings["reporting"]["settings"] # type: ignore[index] + else: + reporting_module = AVAILABLE_REPORTING_MODULES[format] + + report_settings["show_skipped"] = show_skipped + report_settings["output_file"] = output_file + + try: + (modulename, classname) = reporting_module.rsplit(".", 1) + imported_module = importlib.import_module(name=modulename, package=classname) + + report_class: Report = getattr(imported_module, classname) + report_class.generate(scan=modelscan, settings=report_settings) - # ConsoleReport.generate( - # scan=modelscan, - # show_skipped=show_skipped, - # settings=settings["reporting"] - # ) + except Exception as e: + logger.error(f"Error generating report using {reporting_module}: {e}") # exit code 3 if no supported files were passed if not modelscan.scanned: diff --git a/modelscan/reports.py b/modelscan/reports.py index 721e05a..231bf23 100644 --- a/modelscan/reports.py +++ b/modelscan/reports.py @@ -23,7 +23,6 @@ def __init__(self) -> None: @staticmethod def generate( scan: ModelScan, - show_skipped: bool = False, settings: Dict[str, Any] = {}, ) -> Optional[str]: """ @@ -41,7 +40,6 @@ class ConsoleReport(Report): @staticmethod def generate( scan: ModelScan, - show_skipped: bool = False, settings: Dict[str, Any] = {}, ) -> None: issues_by_severity = scan.issues.group_by_severity() @@ -77,7 +75,7 @@ def generate( print( f"\nTotal skipped: {len(scan.skipped)} - run with --show-skipped to see the full list." ) - if show_skipped: + if settings["show_skipped"]: print(f"\nSkipped files list:\n") for file_name in scan.skipped: print(str(file_name)) @@ -87,11 +85,10 @@ class JSONReport(Report): @staticmethod def generate( scan: ModelScan, - show_skipped: bool = False, settings: Dict[str, Any] = {}, ) -> None: report: Dict[str, Any] = scan._generate_results() - if not show_skipped: + if not settings["show_skipped"]: del report["skipped"] print(json.dumps(report)) diff --git a/modelscan/settings.py b/modelscan/settings.py index 00b008d..441adfa 100644 --- a/modelscan/settings.py +++ b/modelscan/settings.py @@ -2,6 +2,10 @@ from typing import Any +AVAILABLE_REPORTING_MODULES = { + "console": "modelscan.reports.ConsoleReport", + "json": "modelscan.reports.JSONReport", +} DEFAULT_SETTINGS = { "supported_zip_extensions": [".zip", ".npz"], @@ -86,11 +90,8 @@ }, "MEDIUM": {}, "LOW": {}, - "reporting_module": { - "module": "modelscan.reports.ConsoleReport", - "settings": {}, - }, }, + "reporting": {"module": "modelscan.reports.ConsoleReport", "settings": {}}, } From d92c83874e51eb91b866cac99734f8fc4aef1a51 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 12 Jan 2024 16:10:06 -0800 Subject: [PATCH 3/5] format help option --- modelscan/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelscan/cli.py b/modelscan/cli.py index 1ffd394..f002559 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -79,7 +79,7 @@ def cli() -> None: "--format", type=click.Choice(["console", "json", "custom"]), default="console", - help="Format of the output. Options are console or json (default: console)", + help="Format of the output. Options are console, json, or custom (to be defined in settings-file). Default is console.", ) @click.option( "-o", From 8c2aab54aca2a0592632263057c78c7f9bbee230 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Tue, 16 Jan 2024 12:16:02 -0800 Subject: [PATCH 4/5] update naming and comment --- modelscan/cli.py | 15 +++++++-------- modelscan/settings.py | 7 +++++-- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/modelscan/cli.py b/modelscan/cli.py index f002559..880063b 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -14,7 +14,7 @@ from modelscan.settings import ( SettingsUtils, DEFAULT_SETTINGS, - AVAILABLE_REPORTING_MODULES, + DEFAULT_REPORTING_MODULES, ) from modelscan.tools.cli_utils import DefaultGroup @@ -75,8 +75,7 @@ def cli() -> None: help="Specify a settings file to use for the scan. Defaults to ./modelscan-settings.toml.", ) @click.option( - "-f", - "--format", + "--reporting-format", type=click.Choice(["console", "json", "custom"]), default="console", help="Format of the output. Options are console, json, or custom (to be defined in settings-file). Default is console.", @@ -86,7 +85,7 @@ def cli() -> None: "--output-file", type=click.Path(), default=None, - help="Optional json reporting output file", + help="Optional file name for output report", ) @cli.command( help="[Default] Scan a model file or diretory for ability to execute suspicious actions. " @@ -98,7 +97,7 @@ def scan( path: Optional[str], show_skipped: bool, settings_file: Optional[str], - format: str, + reporting_format: str, output_file: Path, ) -> int: logger.setLevel(logging.INFO) @@ -134,12 +133,12 @@ def scan( raise click.UsageError("Command line must include a path") report_settings: Dict[str, Any] = {} - if format == "custom": + if reporting_format == "custom": reporting_module = settings["reporting"]["module"] # type: ignore[index] - report_settings = settings["reporting"]["settings"] # type: ignore[index] else: - reporting_module = AVAILABLE_REPORTING_MODULES[format] + 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 diff --git a/modelscan/settings.py b/modelscan/settings.py index 441adfa..c4bad2e 100644 --- a/modelscan/settings.py +++ b/modelscan/settings.py @@ -2,7 +2,7 @@ from typing import Any -AVAILABLE_REPORTING_MODULES = { +DEFAULT_REPORTING_MODULES = { "console": "modelscan.reports.ConsoleReport", "json": "modelscan.reports.JSONReport", } @@ -91,7 +91,10 @@ "MEDIUM": {}, "LOW": {}, }, - "reporting": {"module": "modelscan.reports.ConsoleReport", "settings": {}}, + "reporting": { + "module": "modelscan.reports.ConsoleReport", + "settings": {}, + }, # JSON reporting can be configured by changing "module" to "modelscan.reports.JSONReport" and adding an optional "output_file" field. For custom reporting modules, change "module" to the module name and add the applicable settings fields } From 417842a76d45cd623de562700657d7ef9bf71b0c Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Wed, 17 Jan 2024 09:16:42 -0800 Subject: [PATCH 5/5] -r option --- modelscan/cli.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelscan/cli.py b/modelscan/cli.py index 880063b..db93869 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -75,6 +75,7 @@ def cli() -> None: help="Specify a settings file to use for the scan. Defaults to ./modelscan-settings.toml.", ) @click.option( + "-r", "--reporting-format", type=click.Choice(["console", "json", "custom"]), default="console",