diff --git a/modelscan/modelscan.py b/modelscan/modelscan.py index 4442f5eb..e8308e2f 100644 --- a/modelscan/modelscan.py +++ b/modelscan/modelscan.py @@ -197,7 +197,7 @@ def _scan_source( ) if scan_results.errors: self._errors.extend(scan_results.errors) - elif scan_results.issues: + if scan_results.issues: self._scanned.append(str(model.get_source())) self._issues.add_issues(scan_results.issues) diff --git a/modelscan/tools/picklescanner.py b/modelscan/tools/picklescanner.py index 44c4e2a0..bca8663f 100644 --- a/modelscan/tools/picklescanner.py +++ b/modelscan/tools/picklescanner.py @@ -16,6 +16,12 @@ from .utils import MAGIC_NUMBER, _should_read_directly, get_magic_number +def _source_looks_like_pickle_stream(model: Model) -> bool: + source = str(model.get_source()).lower() + source_leaf = source.split(":")[-1] + return source_leaf.endswith((".pkl", ".pickle", ".joblib", ".dat", ".data", ".npy")) + + class GenOpsError(Exception): def __init__(self, msg: str, globals: Optional[Set[Tuple[str, str]]]): self.msg = msg @@ -57,15 +63,13 @@ def _list_globals( last_byte = b"dummy" while last_byte != b"": # List opcodes + ops: List[Tuple[Any, Any, Union[int, None]]] = [] + parsing_error: Optional[Exception] = None try: - ops: List[Tuple[Any, Any, Union[int, None]]] = list( - pickletools.genops(data) - ) + for op in pickletools.genops(data): + ops.append(op) except Exception as e: - # Given we can have multiple pickles in a file, we may have already successfully extracted globals from a valid pickle. - # Thus return the already found globals in the error & let the caller decide what to do. - globals_opt = globals if len(globals) > 0 else None - raise GenOpsError(str(e), globals_opt) + parsing_error = e last_byte = data.read(1) data.seek(-1, 1) @@ -113,6 +117,12 @@ def _list_globals( f"Found {len(values)} values for STACK_GLOBAL at position {n} instead of 2." ) globals.add((values[1], values[0])) + + if parsing_error is not None: + # A partial parse means we reached executable pickle opcodes before hitting malformed bytes. + # Preserve those globals and let the caller mark the file unsafe. + globals_opt = globals if len(globals) > 0 else None + raise GenOpsError(str(parsing_error), globals_opt) if not multiple_pickles: break @@ -132,11 +142,46 @@ def scan_pickle_bytes( 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( + results = _build_scan_result_from_raw_globals( e.globals, model, settings, ) + if len(results.issues) == 0: + results = _build_scan_result_from_raw_globals( + {("unknown", "pickle_parsing_error")}, + model, + settings, + ) + return ScanResults( + results.issues, + [ + PickleGenopsError( + scan_name, + f"Parsing error: {e}", + model, + ) + ], + [], + ) + if _source_looks_like_pickle_stream(model): + raw_globals_with_parse_error = {("unknown", "pickle_parsing_error")} + results = _build_scan_result_from_raw_globals( + raw_globals_with_parse_error, + model, + settings, + ) + return ScanResults( + results.issues, + [ + PickleGenopsError( + scan_name, + f"Parsing error: {e}", + model, + ) + ], + [], + ) return ScanResults( issues, [ diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index a82e4ecc..67689a27 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -606,6 +606,54 @@ def test_scan_numpy(numpy_file_path: Any) -> None: assert results["errors"] == [] +def test_pickle_parse_error_without_dangerous_globals_fails_closed() -> None: + """Verify hybrid joblib parse errors without pre-error dangerous globals fail closed.""" + # Create a hybrid stream: valid pickle opcodes followed by binary junk + import pickle + import io + + safe_pickle = pickle.dumps(["a", "b", "c"], protocol=4) + hybrid_stream = safe_pickle + b"\xff\xff\xff\xfe\xfd\xfc" # Invalid opcodes + + model = Model("test_hybrid.joblib", io.BytesIO(hybrid_stream)) + results = scan_pickle_bytes(model, settings) + # Should report sentinel unknown.pickle_parsing_error since no real dangerous globals + assert len(results.issues) >= 1 + assert any( + issue.details.operator == "pickle_parsing_error" + for issue in results.issues + ) + + +def test_pickle_parse_error_with_dangerous_globals_no_sentinel() -> None: + """Verify parse errors with already-detected dangerous globals do NOT add duplicate sentinel.""" + import pickle + import io + + class Evil: + def __reduce__(self): + return (eval, ("1+1",)) + + # Dangerous pickle followed by junk that triggers parse error + dangerous_pickle = pickle.dumps(Evil(), protocol=4) + hybrid_stream = dangerous_pickle + b"\xff\xff\xff\xfe\xfd\xfc" + + model = Model("test_dangerous_hybrid.joblib", io.BytesIO(hybrid_stream)) + results = scan_pickle_bytes(model, settings) + # Should report eval, not add duplicate sentinel + assert len(results.issues) >= 1 + assert any( + issue.details.operator == "eval" and issue.details.module == "builtins" + for issue in results.issues + ) + # Should NOT have pickle_parsing_error sentinel if real dangerous globals found + sentinel_count = sum( + 1 for issue in results.issues + if issue.details.operator == "pickle_parsing_error" + ) + assert sentinel_count == 0 + + def test_scan_file_path(file_path: Any) -> None: benign_pickle = ModelScan() results = benign_pickle.scan(Path(f"{file_path}/data/benign0_v3.pkl"))