PyTorch Zipfile Check and Tests (#87)

protectai · Jan 19, 2024 · ee254ab · ee254ab
1 parent 3ea3473
commit ee254ab
Show file tree

Hide file tree

Showing 5 changed files with 424 additions and 49 deletions.
diff --git a/modelscan/scanners/pickle/scan.py b/modelscan/scanners/pickle/scan.py
@@ -3,6 +3,7 @@
 from typing import IO, Union, Optional
 
 from modelscan.scanners.scan import ScanBase, ScanResults
+from modelscan.tools.utils import _is_zipfile
 from modelscan.tools.picklescanner import (
     scan_numpy,
     scan_pickle_bytes,
@@ -26,6 +27,9 @@ def scan(
         ):
             return None
 
+        if _is_zipfile(source):
+            return None
+
         if data:
             results = scan_pytorch(data=data, source=source, settings=self._settings)