From 73e174bec0dc6dea7d876dbe03c4c16bbebff378 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Wed, 16 Aug 2023 14:20:37 -0700 Subject: [PATCH 1/6] update exit codes Return 1 for issues found, 0 for safe --- README.md | 5 +++++ modelscan/cli.py | 10 +++++++--- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index bc2b8a8..f0e899c 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,11 @@ it allows an attacker to read our AWS credentials and write them to another plac That is a firm NO for usage. +The CLI exit status codes are (a-la [ClamAV](https://www.clamav.net/)): +- `0`: scan did not find malware +- `1`: scan found malware +- `2`: scan failed + ## Integrating ModelScan In Your ML Pipelines and CI/CD Pipelines Ad-hoc scanning is a great first step, please drill it into yourself, peers, and friends to do diff --git a/modelscan/cli.py b/modelscan/cli.py index f4396d6..b2df162 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -75,16 +75,20 @@ def cli( "Command line must include either a path or a Hugging Face model" ) ConsoleReport.generate(modelscan.issues, modelscan.errors) - return 0 + + if modelscan.issues.all_issues: + sys.exit(1) + else: + sys.exit(0) except click.UsageError as e: click.echo(e) click.echo(ctx.get_help()) - return 2 + sys.exit(2) except Exception as e: logger.exception(f"Exception: {e}") - return 2 + sys.exit(2) if __name__ == "__main__": From 453dba92bd9c6f41c18e17069a0e348b5ef46207 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 18 Aug 2023 10:48:29 -0700 Subject: [PATCH 2/6] Full 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 --- README.md | 10 +++-- malicious0.pt | 28 ++++++++++++++ modelscan/cli.py | 86 +++++++++++++++++++++++++++++------------- modelscan/modelscan.py | 6 +++ pyproject.toml | 2 +- 5 files changed, 100 insertions(+), 32 deletions(-) create mode 100644 malicious0.pt diff --git a/README.md b/README.md index f0e899c..3ca47fb 100644 --- a/README.md +++ b/README.md @@ -143,10 +143,12 @@ it allows an attacker to read our AWS credentials and write them to another plac That is a firm NO for usage. -The CLI exit status codes are (a-la [ClamAV](https://www.clamav.net/)): -- `0`: scan did not find malware -- `1`: scan found malware -- `2`: scan failed +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 ## Integrating ModelScan In Your ML Pipelines and CI/CD Pipelines diff --git a/malicious0.pt b/malicious0.pt new file mode 100644 index 0000000..05fb0ba --- /dev/null +++ b/malicious0.pt @@ -0,0 +1,28 @@ +c__builtin__ +globals +(tRp100 +0c__builtin__ +compile +(S'fl=open("/etc/passwd");picklesmashed=fl.read();' +S'' +S'exec' +tRp101 +0c__builtin__ +eval +(g101 +g100 +tRp102 +0c__builtin__ +getattr +(c__builtin__ +dict +S'get' +tRp103 +0c__builtin__ +apply +(g103 +(g100 +S'picklesmashed' +ltRp104 +0g104 +. \ No newline at end of file diff --git a/modelscan/cli.py b/modelscan/cli.py index b2df162..8b435a7 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") @@ -58,38 +66,62 @@ 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.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" + ) + ConsoleReport.generate(modelscan.issues, modelscan.errors) - if modelscan.issues.all_issues: - sys.exit(1) - else: - sys.exit(0) + # 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 - except click.UsageError as e: - click.echo(e) - click.echo(ctx.get_help()) - sys.exit(2) + +def main() -> None: + try: + result = cli.main(standalone_mode=False) + + except ( + click.UsageError, + click.BadParameter, + click.BadOptionUsage, + click.NoSuchOption, + click.ClickException, + ) as e: + click.echo(f"Usage 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}") - sys.exit(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 4707107..7991498 100644 --- a/modelscan/modelscan.py +++ b/modelscan/modelscan.py @@ -38,6 +38,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(): @@ -108,6 +109,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) @@ -137,3 +139,7 @@ def issues(self) -> Issues: @property def errors(self) -> List[Error]: return self._errors + + @property + def scanned(self) -> List[str]: + return self._scanned diff --git a/pyproject.toml b/pyproject.toml index fd015f1..c55f497 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" From afe2bb1027da0f89f4d63c0214c2edfe504ff398 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 18 Aug 2023 10:51:28 -0700 Subject: [PATCH 3/6] Remove extra test file --- malicious0.pt | 28 ---------------------------- 1 file changed, 28 deletions(-) delete mode 100644 malicious0.pt diff --git a/malicious0.pt b/malicious0.pt deleted file mode 100644 index 05fb0ba..0000000 --- a/malicious0.pt +++ /dev/null @@ -1,28 +0,0 @@ -c__builtin__ -globals -(tRp100 -0c__builtin__ -compile -(S'fl=open("/etc/passwd");picklesmashed=fl.read();' -S'' -S'exec' -tRp101 -0c__builtin__ -eval -(g101 -g100 -tRp102 -0c__builtin__ -getattr -(c__builtin__ -dict -S'get' -tRp103 -0c__builtin__ -apply -(g103 -(g100 -S'picklesmashed' -ltRp104 -0g104 -. \ No newline at end of file From 64b6c6bd053d9d446503dc51d64c17e1dfcbaf95 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Wed, 27 Sep 2023 10:59:19 -0700 Subject: [PATCH 4/6] add back Modelscan.scanned() --- modelscan/cli.py | 2 +- modelscan/modelscan.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/modelscan/cli.py b/modelscan/cli.py index c436819..4b36a61 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -96,7 +96,7 @@ def cli( ) # exit code 3 if no supported files were passed - if not modelscan._scanned: + if not modelscan.scanned: return 3 # exit code 2 if scan encountered errors elif modelscan.errors: diff --git a/modelscan/modelscan.py b/modelscan/modelscan.py index 0240fc0..77ab30a 100644 --- a/modelscan/modelscan.py +++ b/modelscan/modelscan.py @@ -156,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 From 2f456123f0d21bdd14b227bffe95a3910a217ae4 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Thu, 5 Oct 2023 12:35:53 -0700 Subject: [PATCH 5/6] readme, condense click exceptions --- README.md | 15 ++++++++------- modelscan/cli.py | 10 ++-------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 89e146b..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: @@ -143,13 +151,6 @@ it allows an attacker to read our AWS credentials and write them to another plac That is a firm NO for usage. -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 - ## Integrating ModelScan In Your ML Pipelines and CI/CD Pipelines Ad-hoc scanning is a great first step, please drill it into yourself, peers, and friends to do diff --git a/modelscan/cli.py b/modelscan/cli.py index 4b36a61..5173496 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -113,14 +113,8 @@ def main() -> None: try: result = cli.main(standalone_mode=False) - except ( - click.UsageError, - click.BadParameter, - click.BadOptionUsage, - click.NoSuchOption, - click.ClickException, - ) as e: - click.echo(f"Usage Error: {e}") + 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 From 58022f14e2e5695d9472891d5169aa4a8b8e2068 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Mon, 9 Oct 2023 13:46:10 -0700 Subject: [PATCH 6/6] remove () --- modelscan/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelscan/cli.py b/modelscan/cli.py index 5173496..656980f 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -113,7 +113,7 @@ def main() -> None: try: result = cli.main(standalone_mode=False) - except (click.ClickException,) as e: + except click.ClickException as e: click.echo(f"Error: {e}") with click.Context(cli) as ctx: click.echo(cli.get_help(ctx))