Skip to content

Commit

Permalink
Numpy Header Stream Offset (#164)
Browse files Browse the repository at this point in the history
  • Loading branch information
swashko authored Jun 17, 2024
1 parent a5ca786 commit 27dc4b9
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 4 deletions.
4 changes: 2 additions & 2 deletions modelscan/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
5 changes: 3 additions & 2 deletions modelscan/tools/picklescanner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 27dc4b9

Please sign in to comment.