Numpy Header Stream Offset (#164)

Author: swashko

modelscan/model.py

@@ -47,9 +47,9 @@ def __exit__(self, exc_type, exc_value, traceback) -> None:  # type: ignore
     def get_source(self) -> Path:
         return self._source
 
-    def get_stream(self) -> IO[bytes]:
+    def get_stream(self, offset: int = 0) -> IO[bytes]:
         if not self._stream:
             raise ModelDataEmpty("Model data is empty.")
 
-        self._stream.seek(0)
+        self._stream.seek(offset)
         return self._stream

modelscan/tools/picklescanner.py

@@ -124,11 +124,12 @@ def scan_pickle_bytes(
     settings: Dict[str, Any],
     scan_name: str = "pickle",
     multiple_pickles: bool = True,
+    offset: int = 0,
 ) -> ScanResults:
     """Disassemble a Pickle stream and report issues"""
     issues: List[Issue] = []
     try:
-        raw_globals = _list_globals(model.get_stream(), multiple_pickles)
+        raw_globals = _list_globals(model.get_stream(offset), multiple_pickles)
     except GenOpsError as e:
         if e.globals is not None:
             return _build_scan_result_from_raw_globals(
@@ -231,7 +232,7 @@ def scan_numpy(model: Model, settings: Dict[str, Any]) -> ScanResults:
         _, _, dtype = np.lib.format._read_array_header(stream, version)  # type: ignore[attr-defined]
 
         if dtype.hasobject:
-            return scan_pickle_bytes(model, settings, scan_name)
+            return scan_pickle_bytes(model, settings, scan_name, True, stream.tell())
         else:
             return ScanResults([], [], [])
     else: