From 3c87e3de0fda83902e6f73aa9c80cd2cf1fb99d1 Mon Sep 17 00:00:00 2001
From: Sam Washko <sam@protectai.com>
Date: Mon, 9 Oct 2023 13:48:55 -0700
Subject: [PATCH] Updated Exit Codes (#42)

---
 README.md              |  8 ++++
 modelscan/cli.py       | 84 ++++++++++++++++++++++++++++--------------
 modelscan/modelscan.py |  6 +++
 pyproject.toml         |  2 +-
 4 files changed, 72 insertions(+), 28 deletions(-)

diff --git a/README.md b/README.md
index 3c7a93d..1e298d7 100644
--- a/README.md
+++ b/README.md
@@ -130,6 +130,14 @@ Remember models are just like any other form of digital media, you should scan c
 **NOTE**: LLMs are large files, it can take a few minutes to download them before scanning. Expect the process
 to take just a few minutes to complete. 
 
+##### CLI Exit Codes
+The CLI exit status codes are:
+- `0`: Scan completed successfully, no vulnerabilities found
+- `1`: Scan completed successfully, vulnerabilities found
+- `2`: Scan failed, modelscan threw an error while scanning
+- `3`: No supported files were passed to the tool
+- `4`: Usage error, CLI was passed invalid or incomplete options
+
 ### Understanding The Results
 
 Once a scan has been completed you'll see output like this if an issue is found:
diff --git a/modelscan/cli.py b/modelscan/cli.py
index c83fc03..656980f 100644
--- a/modelscan/cli.py
+++ b/modelscan/cli.py
@@ -15,8 +15,16 @@
 CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
 
 
+# redefine format_usage so the appropriate command name shows up
+class ModelscanCommand(click.Command):
+    def format_usage(self, ctx: click.Context, formatter: click.HelpFormatter) -> None:
+        pieces = self.collect_usage_pieces(ctx)
+        formatter.write_usage("modelscan", " ".join(pieces))
+
+
 @click.command(
     context_settings=CONTEXT_SETTINGS,
+    cls=ModelscanCommand,
     help="Modelscan detects machine learning model files that perform suspicious actions",
 )
 @click.version_option(__version__, "-v", "--version")
@@ -65,39 +73,61 @@ def cli(
     if log is not None:
         logger.setLevel(getattr(logging, log))
 
-    try:
-        modelscan = Modelscan()
-        if path is not None:
-            pathlibPath = Path().cwd() if path == "." else Path(path).absolute()
-            if not pathlibPath.exists():
-                raise FileNotFoundError(f"Path {path} does not exist")
-            else:
-                modelscan.scan_path(pathlibPath)
-        # elif url is not None:
-        #     modelscan.scan_url(url)
-        elif huggingface is not None:
-            modelscan.scan_huggingface_model(huggingface)
+    modelscan = Modelscan()
+    if path is not None:
+        pathlibPath = Path().cwd() if path == "." else Path(path).absolute()
+        if not pathlibPath.exists():
+            raise FileNotFoundError(f"Path {path} does not exist")
         else:
-            raise click.UsageError(
-                "Command line must include either a path or a Hugging Face model"
-            )
-        ConsoleReport.generate(
-            modelscan.issues,
-            modelscan.errors,
-            modelscan._skipped,
-            show_skipped=show_skipped,
+            modelscan.scan_path(pathlibPath)
+    # elif url is not None:
+    #     modelscan.scan_url(url)
+    elif huggingface is not None:
+        modelscan.scan_huggingface_model(huggingface)
+    else:
+        raise click.UsageError(
+            "Command line must include either a path or a Hugging Face model"
         )
-        return 0
+    ConsoleReport.generate(
+        modelscan.issues,
+        modelscan.errors,
+        modelscan._skipped,
+        show_skipped=show_skipped,
+    )
 
-    except click.UsageError as e:
-        click.echo(e)
-        click.echo(ctx.get_help())
+    # exit code 3 if no supported files were passed
+    if not modelscan.scanned:
+        return 3
+    # exit code 2 if scan encountered errors
+    elif modelscan.errors:
         return 2
+    # exit code 1 if scan completed successfully and vulnerabilities were found
+    elif modelscan.issues.all_issues:
+        return 1
+    # exit code 0 if scan completed successfully and no vulnerabilities were found
+    else:
+        return 0
+
+
+def main() -> None:
+    try:
+        result = cli.main(standalone_mode=False)
+
+    except click.ClickException as e:
+        click.echo(f"Error: {e}")
+        with click.Context(cli) as ctx:
+            click.echo(cli.get_help(ctx))
+        # exit code 4 for CLI usage errors
+        result = 4
 
     except Exception as e:
-        logger.exception(f"Exception: {e}")
-        return 2
+        click.echo(f"Exception: {e}")
+        # exit code 2 if scan throws exceptions
+        result = 2
+
+    finally:
+        sys.exit(result)
 
 
 if __name__ == "__main__":
-    sys.exit(cli())
+    main()
diff --git a/modelscan/modelscan.py b/modelscan/modelscan.py
index ced1354..77ab30a 100644
--- a/modelscan/modelscan.py
+++ b/modelscan/modelscan.py
@@ -42,6 +42,7 @@ def __init__(self) -> None:
         self._issues = Issues()
         self._errors: List[Error] = []
         self._skipped: List[str] = []
+        self._scanned: List[str] = []
 
     def scan_path(self, path: Path) -> None:
         if path.is_dir():
@@ -124,6 +125,7 @@ def _scan_source(
             if extension in scan.supported_extensions():
                 logger.info(f"Scanning {source} using {scan.name()} model scan")
                 issues, errors = scan.scan(source=source, data=data)
+                self._scanned.append(str(source))
 
         self._issues.add_issues(issues)
         self._errors.extend(errors)
@@ -154,6 +156,10 @@ def issues(self) -> Issues:
     def errors(self) -> List[Error]:
         return self._errors
 
+    @property
+    def scanned(self) -> List[str]:
+        return self._scanned
+
     @property
     def skipped(self) -> List[str]:
         return self._skipped
diff --git a/pyproject.toml b/pyproject.toml
index d73efcd..f55bf7b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,7 +9,7 @@ packages = [{ include = "modelscan" }]
 exclude = ["tests/*", "Makefile"]
 
 [tool.poetry.scripts]
-modelscan = "modelscan.cli:cli"
+modelscan = "modelscan.cli:main"
 
 [tool.poetry.dependencies]
 python = ">=3.8,<3.12"