From 3777e0b8745f248fff87b17bdd9506de968026a4 Mon Sep 17 00:00:00 2001
From: Sam Washko <sam@protectai.com>
Date: Wed, 13 Mar 2024 16:46:50 -0700
Subject: [PATCH] scan error handling

---
 modelscan/modelscan.py            | 22 ++++++++++++++++++----
 modelscan/scanners/pickle/scan.py |  7 +++++--
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/modelscan/modelscan.py b/modelscan/modelscan.py
index 7627fac..325df14 100644
--- a/modelscan/modelscan.py
+++ b/modelscan/modelscan.py
@@ -122,10 +122,24 @@ def _scan_source(
         scanned = False
         for scan_class in self._scanners_to_run:
             scanner = scan_class(self._settings)  # type: ignore[operator]
-            scan_results = scanner.scan(
-                source=source,
-                data=data,
-            )
+
+            try:
+                scan_results = scanner.scan(
+                    source=source,
+                    data=data,
+                )
+            except Exception as e:
+                logger.error(
+                    f"Error encountered from scanner {scanner.full_name()} with path {source}: {e}"
+                )
+                self._errors.append(
+                    ModelScanError(
+                        scanner.full_name(),
+                        ErrorCategories.MODEL_SCAN,
+                        f"Error encountered from scanner {scanner.full_name()}: {e}",
+                        f"{source}",
+                    )
+                )
 
             if scan_results is not None:
                 scanned = True
diff --git a/modelscan/scanners/pickle/scan.py b/modelscan/scanners/pickle/scan.py
index 2c793ae..feb8272 100644
--- a/modelscan/scanners/pickle/scan.py
+++ b/modelscan/scanners/pickle/scan.py
@@ -67,8 +67,11 @@ def scan(
         if data:
             results = scan_numpy(data=data, source=source, settings=self._settings)
 
-        with open(source, "rb") as file_io:
-            results = scan_numpy(data=file_io, source=source, settings=self._settings)
+        else:
+            with open(source, "rb") as file_io:
+                results = scan_numpy(
+                    data=file_io, source=source, settings=self._settings
+                )
 
         return self.label_results(results)