From 6b7c15164c082380531cd0dc15cfb555e0f752de Mon Sep 17 00:00:00 2001 From: Dream <2468001320@qq.com> Date: Thu, 27 Aug 2026 17:23:43 +0800 Subject: [PATCH 1/2] build: make NumPy an optional dependency - pyproject.toml --- pyproject.toml | 121 +++++++++++++++++++++++++------------------------ 1 file changed, 61 insertions(+), 60 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f538120..e0a4766 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,60 +1,61 @@ -[tool.poetry] -name = "modelscan" -version = "0.0.0" -description = "The modelscan package is a cli tool for detecting unsafe operations in model files across various model serialization formats." -authors = ["ProtectAI "] -license = "Apache License 2.0" -readme = "README.md" -packages = [{ include = "modelscan" }] -exclude = ["tests/*", "Makefile"] - -[tool.poetry.scripts] -modelscan = "modelscan.cli:main" - -[tool.poetry.dependencies] -python = ">=3.10,<3.13" -click = "^8.1.3" -numpy = ">=1.24.3" -rich = ">=13.4.2,<15.0.0" -tomlkit = ">=0.12.3,<0.14.0" -h5py = { version = "^3.9.0", optional = true } - -# TODO: Add py3.12 once TF release supports -tensorflow = { version = "^2.17", optional = true } - -[tool.poetry.extras] -tensorflow = ["tensorflow"] -h5py = ["h5py"] - -[tool.poetry.group.test.dependencies] -pytest = ">=7.4,<9.0" -bandit = { version = "1.8.3", extras = ["toml"] } -mypy = "^1.4.1" -requests = "^2.31.0" -aiohttp = "^3.8.5" -dill = ">=0.3.7,<0.5.0" -types-requests = ">1.26" -torch = "^2.7.0" -tf-keras = "^2.20.1" - - -[tool.poetry.group.dev.dependencies] -dunamai = "^1.18.0" -pre-commit = ">=3.3.3,<5.0.0" -black = ">=23.7,<26.0" - -[tool.poetry.group.prod.dependencies] -dunamai = "^1.18.0" - -[build-system] -requires = ["poetry-core", "poetry-dynamic-versioning"] -build-backend = "poetry.core.masonry.api" - -[tool.poetry-dynamic-versioning] -enable = true - -[tool.bandit] -exclude_dirs = ["tests", "notebooks"] - -[tool.mypy] -exclude = ["notebooks"] +[tool.poetry] +name = "modelscan" +version = "0.0.0" +description = "The modelscan package is a cli tool for detecting unsafe operations in model files across various model serialization formats." +authors = ["ProtectAI "] +license = "Apache License 2.0" +readme = "README.md" +packages = [{ include = "modelscan" }] +exclude = ["tests/*", "Makefile"] + +[tool.poetry.scripts] +modelscan = "modelscan.cli:main" + +[tool.poetry.dependencies] +python = ">=3.10,<3.13" +click = "^8.1.3" +numpy = { version = ">=1.24.3", optional = true } +rich = ">=13.4.2,<15.0.0" +tomlkit = ">=0.12.3,<0.14.0" +h5py = { version = "^3.9.0", optional = true } + +# TODO: Add py3.12 once TF release supports +tensorflow = { version = "^2.17", optional = true } + +[tool.poetry.extras] +tensorflow = ["tensorflow"] +h5py = ["h5py"] +numpy = ["numpy"] + +[tool.poetry.group.test.dependencies] +pytest = ">=7.4,<9.0" +bandit = { version = "1.8.3", extras = ["toml"] } +mypy = "^1.4.1" +requests = "^2.31.0" +aiohttp = "^3.8.5" +dill = ">=0.3.7,<0.5.0" +types-requests = ">1.26" +torch = "^2.7.0" +tf-keras = "^2.20.1" + + +[tool.poetry.group.dev.dependencies] +dunamai = "^1.18.0" +pre-commit = ">=3.3.3,<5.0.0" +black = ">=23.7,<26.0" + +[tool.poetry.group.prod.dependencies] +dunamai = "^1.18.0" + +[build-system] +requires = ["poetry-core", "poetry-dynamic-versioning"] +build-backend = "poetry.core.masonry.api" + +[tool.poetry-dynamic-versioning] +enable = true + +[tool.bandit] +exclude_dirs = ["tests", "notebooks"] + +[tool.mypy] +exclude = ["notebooks"] From 7d2439c252f19214207e75432947b0a59ebcb37e Mon Sep 17 00:00:00 2001 From: Dream <2468001320@qq.com> Date: Thu, 27 Aug 2026 17:23:47 +0800 Subject: [PATCH 2/2] build: make NumPy an optional dependency - modelscan/tools/picklescanner.py --- modelscan/tools/picklescanner.py | 546 ++++++++++++++++--------------- 1 file changed, 277 insertions(+), 269 deletions(-) diff --git a/modelscan/tools/picklescanner.py b/modelscan/tools/picklescanner.py index 44c4e2a..8377cc0 100644 --- a/modelscan/tools/picklescanner.py +++ b/modelscan/tools/picklescanner.py @@ -1,269 +1,277 @@ -import logging -import pickletools # nosec -from tarfile import TarError -from typing import IO, Any, Dict, List, Set, Tuple, Union, Optional - -import numpy as np - -from modelscan.error import PickleGenopsError -from modelscan.skip import ModelScanSkipped, SkipCategories -from modelscan.issues import Issue, IssueCode, IssueSeverity, OperatorIssueDetails -from modelscan.scanners.scan import ScanResults -from modelscan.model import Model - -logger = logging.getLogger("modelscan") - -from .utils import MAGIC_NUMBER, _should_read_directly, get_magic_number - - -class GenOpsError(Exception): - def __init__(self, msg: str, globals: Optional[Set[Tuple[str, str]]]): - self.msg = msg - self.globals = globals - super().__init__() - - def __str__(self) -> str: - return self.msg - - -# TODO: handle methods loading other Pickle files (either mark as suspicious, or follow calls to scan other files [preventing infinite loops]) -# -# pickle.loads() -# https://docs.python.org/3/library/pickle.html#pickle.loads -# pickle.load() -# https://docs.python.org/3/library/pickle.html#pickle.load -# numpy.load() -# https://numpy.org/doc/stable/reference/generated/numpy.load.html#numpy.load -# numpy.ctypeslib.load_library() -# https://numpy.org/doc/stable/reference/routines.ctypeslib.html#numpy.ctypeslib.load_library -# pandas.read_pickle() -# https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_pickle.html -# joblib.load() -# https://joblib.readthedocs.io/en/latest/generated/joblib.load.html -# torch.load() -# https://pytorch.org/docs/stable/generated/torch.load.html -# tf.keras.models.load_model() -# https://www.tensorflow.org/api_docs/python/tf/keras/models/load_model -# - - -def _list_globals( - data: IO[bytes], multiple_pickles: bool = True -) -> Set[Tuple[str, str]]: - globals: Set[Any] = set() - - memo: Dict[Union[int, str], str] = {} - # Scan the data for pickle buffers, stopping when parsing fails or stops making progress - last_byte = b"dummy" - while last_byte != b"": - # List opcodes - try: - ops: List[Tuple[Any, Any, Union[int, None]]] = list( - pickletools.genops(data) - ) - 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) - - last_byte = data.read(1) - data.seek(-1, 1) - - # Extract global imports - for n in range(len(ops)): - op = ops[n] - op_name = op[0].name - op_value: str = op[1] - - if op_name == "MEMOIZE" and n > 0: - memo[len(memo)] = ops[n - 1][1] - elif op_name in ["PUT", "BINPUT", "LONG_BINPUT"] and n > 0: - memo[op_value] = ops[n - 1][1] - elif op_name in ("GLOBAL", "INST"): - globals.add(tuple(op_value.split(" ", 1))) - elif op_name == "STACK_GLOBAL": - values: List[str] = [] - for offset in range(1, n): - if ops[n - offset][0].name in [ - "MEMOIZE", - "PUT", - "BINPUT", - "LONG_BINPUT", - ]: - continue - if ops[n - offset][0].name in ["GET", "BINGET", "LONG_BINGET"]: - values.append(memo[int(ops[n - offset][1])]) - elif ops[n - offset][0].name not in [ - "SHORT_BINUNICODE", - "UNICODE", - "BINUNICODE", - "BINUNICODE8", - ]: - logger.debug( - "Presence of non-string opcode, categorizing as an unknown dangerous import" - ) - values.append("unknown") - else: - values.append(ops[n - offset][1]) - if len(values) == 2: - break - if len(values) != 2: - raise ValueError( - f"Found {len(values)} values for STACK_GLOBAL at position {n} instead of 2." - ) - globals.add((values[1], values[0])) - if not multiple_pickles: - break - - return globals - - -def scan_pickle_bytes( - model: Model, - 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(offset), multiple_pickles) - except GenOpsError as e: - if e.globals is not None: - return _build_scan_result_from_raw_globals( - e.globals, - model, - settings, - ) - return ScanResults( - issues, - [ - PickleGenopsError( - scan_name, - f"Parsing error: {e}", - model, - ) - ], - [], - ) - logger.debug("Global imports in %s: %s", model, raw_globals, settings) - return _build_scan_result_from_raw_globals(raw_globals, model, settings) - - -def _build_scan_result_from_raw_globals( - raw_globals: Set[Tuple[str, str]], - model: Model, - settings: Dict[str, Any], -) -> ScanResults: - issues: List[Issue] = [] - severities = { - "CRITICAL": IssueSeverity.CRITICAL, - "HIGH": IssueSeverity.HIGH, - "MEDIUM": IssueSeverity.MEDIUM, - "LOW": IssueSeverity.LOW, - } - - for rg in raw_globals: - global_module, global_name, severity = rg[0], rg[1], None - for severity_name in severities: - if global_module not in settings["unsafe_globals"][severity_name]: - continue - filter = settings["unsafe_globals"][severity_name][global_module] - if filter == "*": - severity = severities[severity_name] - break - for filter_value in filter: - if filter_value in global_name: - severity = severities[severity_name] - break - else: - continue - break - if "unknown" in global_module or "unknown" in global_name: - severity = IssueSeverity.CRITICAL # we must assume it is RCE - if severity is not None: - issues.append( - Issue( - code=IssueCode.UNSAFE_OPERATOR, - severity=severity, - details=OperatorIssueDetails( - module=global_module, - operator=global_name, - source=model.get_source(), - severity=severity, - ), - ) - ) - return ScanResults(issues, [], []) - - -def scan_numpy(model: Model, settings: Dict[str, Any]) -> ScanResults: - scan_name = "numpy" - # Code to distinguish from NumPy binary files and pickles. - _ZIP_PREFIX = b"PK\x03\x04" - _ZIP_SUFFIX = b"PK\x05\x06" # empty zip files start with this - N = len(np.lib.format.MAGIC_PREFIX) - stream = model.get_stream() - magic = stream.read(N) - # If the file size is less than N, we need to make sure not - # to seek past the beginning of the file - stream.seek(-min(N, len(magic)), 1) # back-up - if magic.startswith(_ZIP_PREFIX) or magic.startswith(_ZIP_SUFFIX): - # .npz file - return ScanResults( - [], - [], - [ - ModelScanSkipped( - scan_name, - SkipCategories.NOT_IMPLEMENTED, - "Scanning of .npz files is not implemented yet", - str(model.get_source()), - ) - ], - ) - - elif magic == np.lib.format.MAGIC_PREFIX: - # .npy file - version = np.lib.format.read_magic(stream) # type: ignore[no-untyped-call] - np.lib.format._check_version(version) # type: ignore[attr-defined] - _, _, dtype = np.lib.format._read_array_header(stream, version) # type: ignore[attr-defined] - - if dtype.hasobject: - return scan_pickle_bytes(model, settings, scan_name, True, stream.tell()) - else: - return ScanResults([], [], []) - else: - return scan_pickle_bytes(model, settings, scan_name) - - -def scan_pytorch(model: Model, settings: Dict[str, Any]) -> ScanResults: - scan_name = "pytorch" - should_read_directly = _should_read_directly(model.get_stream()) - if should_read_directly and model.get_stream().tell() == 0: - # try loading from tar - try: - # TODO: implement loading from tar - raise TarError() - except TarError: - # file does not contain a tar - model.get_stream().seek(0) - - magic = get_magic_number(model.get_stream()) - if magic != MAGIC_NUMBER: - return ScanResults( - [], - [], - [ - ModelScanSkipped( - scan_name, - SkipCategories.MAGIC_NUMBER, - "Invalid magic number", - str(model.get_source()), - ) - ], - ) - - return scan_pickle_bytes(model, settings, scan_name, multiple_pickles=False) +import logging +import pickletools # nosec +from tarfile import TarError +from typing import IO, Any, Dict, List, Set, Tuple, Union, Optional + +from modelscan.error import PickleGenopsError +from modelscan.skip import ModelScanSkipped, SkipCategories +from modelscan.issues import Issue, IssueCode, IssueSeverity, OperatorIssueDetails +from modelscan.scanners.scan import ScanResults +from modelscan.model import Model + +logger = logging.getLogger("modelscan") + +from .utils import MAGIC_NUMBER, _should_read_directly, get_magic_number + + +class GenOpsError(Exception): + def __init__(self, msg: str, globals: Optional[Set[Tuple[str, str]]]): + self.msg = msg + self.globals = globals + super().__init__() + + def __str__(self) -> str: + return self.msg + + +# TODO: handle methods loading other Pickle files (either mark as suspicious, or follow calls to scan other files [preventing infinite loops]) +# +# pickle.loads() +# https://docs.python.org/3/library/pickle.html#pickle.loads +# pickle.load() +# https://docs.python.org/3/library/pickle.html#pickle.load +# numpy.load() +# https://numpy.org/doc/stable/reference/generated/numpy.load.html#numpy.load +# numpy.ctypeslib.load_library() +# https://numpy.org/doc/stable/reference/routines.ctypeslib.html#numpy.ctypeslib.load_library +# pandas.read_pickle() +# https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_pickle.html +# joblib.load() +# https://joblib.readthedocs.io/en/latest/generated/joblib.load.html +# torch.load() +# https://pytorch.org/docs/stable/generated/torch.load.html +# tf.keras.models.load_model() +# https://www.tensorflow.org/api_docs/python/tf/keras/models/load_model +# + + +def _list_globals( + data: IO[bytes], multiple_pickles: bool = True +) -> Set[Tuple[str, str]]: + globals: Set[Any] = set() + + memo: Dict[Union[int, str], str] = {} + # Scan the data for pickle buffers, stopping when parsing fails or stops making progress + last_byte = b"dummy" + while last_byte != b"": + # List opcodes + try: + ops: List[Tuple[Any, Any, Union[int, None]]] = list( + pickletools.genops(data) + ) + 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) + + last_byte = data.read(1) + data.seek(-1, 1) + + # Extract global imports + for n in range(len(ops)): + op = ops[n] + op_name = op[0].name + op_value: str = op[1] + + if op_name == "MEMOIZE" and n > 0: + memo[len(memo)] = ops[n - 1][1] + elif op_name in ["PUT", "BINPUT", "LONG_BINPUT"] and n > 0: + memo[op_value] = ops[n - 1][1] + elif op_name in ("GLOBAL", "INST"): + globals.add(tuple(op_value.split(" ", 1))) + elif op_name == "STACK_GLOBAL": + values: List[str] = [] + for offset in range(1, n): + if ops[n - offset][0].name in [ + "MEMOIZE", + "PUT", + "BINPUT", + "LONG_BINPUT", + ]: + continue + if ops[n - offset][0].name in ["GET", "BINGET", "LONG_BINGET"]: + values.append(memo[int(ops[n - offset][1])]) + elif ops[n - offset][0].name not in [ + "SHORT_BINUNICODE", + "UNICODE", + "BINUNICODE", + "BINUNICODE8", + ]: + logger.debug( + "Presence of non-string opcode, categorizing as an unknown dangerous import" + ) + values.append("unknown") + else: + values.append(ops[n - offset][1]) + if len(values) == 2: + break + if len(values) != 2: + raise ValueError( + f"Found {len(values)} values for STACK_GLOBAL at position {n} instead of 2." + ) + globals.add((values[1], values[0])) + if not multiple_pickles: + break + + return globals + + +def scan_pickle_bytes( + model: Model, + 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(offset), multiple_pickles) + except GenOpsError as e: + if e.globals is not None: + return _build_scan_result_from_raw_globals( + e.globals, + model, + settings, + ) + return ScanResults( + issues, + [ + PickleGenopsError( + scan_name, + f"Parsing error: {e}", + model, + ) + ], + [], + ) + logger.debug("Global imports in %s: %s", model, raw_globals, settings) + return _build_scan_result_from_raw_globals(raw_globals, model, settings) + + +def _build_scan_result_from_raw_globals( + raw_globals: Set[Tuple[str, str]], + model: Model, + settings: Dict[str, Any], +) -> ScanResults: + issues: List[Issue] = [] + severities = { + "CRITICAL": IssueSeverity.CRITICAL, + "HIGH": IssueSeverity.HIGH, + "MEDIUM": IssueSeverity.MEDIUM, + "LOW": IssueSeverity.LOW, + } + + for rg in raw_globals: + global_module, global_name, severity = rg[0], rg[1], None + for severity_name in severities: + if global_module not in settings["unsafe_globals"][severity_name]: + continue + filter = settings["unsafe_globals"][severity_name][global_module] + if filter == "*": + severity = severities[severity_name] + break + for filter_value in filter: + if filter_value in global_name: + severity = severities[severity_name] + break + else: + continue + break + if "unknown" in global_module or "unknown" in global_name: + severity = IssueSeverity.CRITICAL # we must assume it is RCE + if severity is not None: + issues.append( + Issue( + code=IssueCode.UNSAFE_OPERATOR, + severity=severity, + details=OperatorIssueDetails( + module=global_module, + operator=global_name, + source=model.get_source(), + severity=severity, + ), + ) + ) + return ScanResults(issues, [], []) + + +def scan_numpy(model: Model, settings: Dict[str, Any]) -> ScanResults: + scan_name = "numpy" + # NumPy is an optional dependency (only needed for .npy scanning); + # import lazily so the package works without it installed. + try: + import numpy as np # noqa: F401 - used below for extension checks + except ImportError as exc: + raise ModelScanSkipped( + model, + SkipCategories.UNSUPPORTED, + "numpy is not installed; skipping NumPy file scanning", + ) from exc + # Code to distinguish from NumPy binary files and pickles. + _ZIP_PREFIX = b"PK\x03\x04" + _ZIP_SUFFIX = b"PK\x05\x06" # empty zip files start with this + N = len(np.lib.format.MAGIC_PREFIX) + stream = model.get_stream() + magic = stream.read(N) + # If the file size is less than N, we need to make sure not + # to seek past the beginning of the file + stream.seek(-min(N, len(magic)), 1) # back-up + if magic.startswith(_ZIP_PREFIX) or magic.startswith(_ZIP_SUFFIX): + # .npz file + return ScanResults( + [], + [], + [ + ModelScanSkipped( + scan_name, + SkipCategories.NOT_IMPLEMENTED, + "Scanning of .npz files is not implemented yet", + str(model.get_source()), + ) + ], + ) + + elif magic == np.lib.format.MAGIC_PREFIX: + # .npy file + version = np.lib.format.read_magic(stream) # type: ignore[no-untyped-call] + np.lib.format._check_version(version) # type: ignore[attr-defined] + _, _, dtype = np.lib.format._read_array_header(stream, version) # type: ignore[attr-defined] + + if dtype.hasobject: + return scan_pickle_bytes(model, settings, scan_name, True, stream.tell()) + else: + return ScanResults([], [], []) + else: + return scan_pickle_bytes(model, settings, scan_name) + + +def scan_pytorch(model: Model, settings: Dict[str, Any]) -> ScanResults: + scan_name = "pytorch" + should_read_directly = _should_read_directly(model.get_stream()) + if should_read_directly and model.get_stream().tell() == 0: + # try loading from tar + try: + # TODO: implement loading from tar + raise TarError() + except TarError: + # file does not contain a tar + model.get_stream().seek(0) + + magic = get_magic_number(model.get_stream()) + if magic != MAGIC_NUMBER: + return ScanResults( + [], + [], + [ + ModelScanSkipped( + scan_name, + SkipCategories.MAGIC_NUMBER, + "Invalid magic number", + str(model.get_source()), + ) + ], + ) + + return scan_pickle_bytes(model, settings, scan_name, multiple_pickles=False)