DOCSEARCHbyAlgolia
Submit Docs

ModelScan

repository·main·Indexed 20 days ago

https://github.com/protectai/modelscan

An open-source CLI tool and Python library from Protect AI designed to detect unsafe code and malicious operations embedded in machine learning model files. It protects against model serialization attacks by scanning formats such as Pickle, H5, and SavedModel across frameworks including PyTorch, TensorFlow, Keras, Sklearn, and XGBoost. The tool identifies vulnerabilities like arbitrary code execution and credential theft without executing the model.

Tokens
10.7K
Snippets
37
Records
57
Agent score
73%

What's inside modelscan

  1. What is a Model Serialization Attack?

    main

    A Model Serialization Attack occurs when malicious code is embedded into a model during the serialization (saving) process. When a user loads the model (e.g., using torch.load(PATH) in PyTorch), the exploit executes immediately upon loading.

    These attacks can be used for:

    • Credential Theft: Stealing cloud credentials to access other systems.
    • Data Theft: Intercepting requests sent to the model.
    • Data Poisoning: Altering data sent after model tasks are performed.
    • Model Poisoning: Altering the actual results/outputs of the model.

  2. Compare ML Model Storage Formats and Security Risks

    main

    Different ML frameworks use different serialization approaches, each with varying levels of vulnerability to code injection:

    ApproachPopularityRisk of ExploitabilityDetails
    Pickle VariantsVery highVery highIncludes pickle, cloudpickle, dill, and joblib. Allows arbitrary code execution.
    TensorFlow SavedModelHighMediumBased on Protocol Buffers. Generally secure, but certain operations like io.read_file or custom operators can be exploited.
    H5 (Keras)HighLowGenerally secure, but Keras Lambda layers allow arbitrary code execution and represent a significant attack surface.
    Inference OnlyMediumLowIncludes ONNX, TensorRT, and Apache TVM. These use internal computational graphs with restricted operators, making them highly secure.
    Vector/Tensor OnlyLowVery lowIncludes Safetensors, JSON, MsgPack, and NPY/NPZ. These store weights/biases without architecture, preventing arbitrary code execution (though weight poisoning is possible).

  3. Supported Model Formats and Frameworks

    main

    ModelScan is designed to protect against attacks across multiple formats and frameworks.

    Supported Formats:

    • H5
    • Pickle
    • SavedModel

    Supported Frameworks (via these formats):

    • PyTorch
    • TensorFlow
    • Keras
    • Sklearn
    • XGBoost

  4. Understand Model Serialization Attacks

    main

    A Model Serialization Attack (also known as a code injection attack) occurs when malicious code is embedded into an ML model during the serialization process (saving to disk). When a user or system loads the model for training or inference, the injected code executes immediately, often without changing the model's visible behavior.

    This attack targets the way models are stored and retrieved, exploiting formats that allow code to be stored alongside model data (vectors/tensors).

  5. Securely Serialize PyTorch Models using ONNX

    main

    To avoid the security risks associated with PyTorch's built-in torch.save (which uses Pickle), you should export your models to the ONNX format. ONNX is a secure format that does not allow arbitrary code execution and is often faster for inference.

    Use the torch.onnx.export method to perform the conversion.

    # Example concept for exporting to ONNX
torch.onnx.export(model, dummy_input, "model.onnx")

  6. Securely Serialize TensorFlow Models

    main

    The native TensorFlow SavedModel format (based on Protocol Buffers) is a recommended secure approach. However, be aware that certain TensorFlow operations can still be exploited for serialization attacks.

    ModelScan is designed to detect these risky operations (such as io.read_file, io.write_file, or io.MatchingFiles) and generate findings to help you secure your models.

  7. Install ModelScan

    main

    ModelScan supports Python 3.9 to 3.12. You can install the base package via pip:

    pip install modelscan

    To include it in your project dependencies, add it to your requirements.txt or pyproject.toml:

    modelscan = ">=0.1.1"

    Note: If you need to scan Tensorflow or HD5 formatted models, you must install the required extras:

    pip install 'modelscan[ tensorflow, h5py ]'

  8. Explore Model Serialization Attack Demonstrations

    main
    The notebooks/ directory contains Jupyter notebooks that demonstrate how model serialization attacks (specifically stealthy mock exfiltration attacks) work across different machine learning libraries. These notebooks show how an attacker can embed malicious code that exfiltrates sensitive data (like AWS secrets) while ensuring the model still functions normally. Each notebook also demonstrates how to use modelscan to detect these unsafe models and provides expected scan results for both safe and unsafe versions.

  9. Securely Serialize Keras Models

    main

    When using Keras, you can use the HDF5 format or the newer Keras v3 format. While these are generally secure, you must avoid using Keras Lambda layers, as they allow arbitrary code execution. Instead, share pre/post-processing code explicitly between your training and inference environments.

    Recommended formats:

    • tf.keras.Model.save with save_format='h5'
    • Keras v3 format (passing save_format='tf' or a filename without an extension)

  10. Best practices for securing ML models

    main

    To protect against model serialization attacks and unauthorized access, implement a defense-in-depth strategy using the following measures:

    1. Authenticated Access: Only store models in systems with authentication (e.g., ensure MLflow instances are behind an authentication gateway).
    2. Least Privilege: Use Authorization or IAM (Identity and Access Management) systems to implement fine-grained access control.
    3. Automated Scanning: Use a tool like ModelScan to catch code injection attempts. Scan models at every stage of the ML ecosystem: before retraining, fine-tuning, evaluation, or inference.
    4. Encryption at Rest: Encrypt model files (e.g., using S3 bucket encryption) to prevent unauthorized reading or writing.
    5. Encryption in Transit: Use TLS or mTLS for all network connections when loading models to protect against Man-in-the-Middle (MITM) attacks.
    6. Integrity Verification: Store and verify checksums for your models to ensure file integrity.
    7. Authenticity: Use cryptographic signatures to ensure both the integrity and the authenticity of the model.

  11. How ModelScan works: Scanners, Middlewares, and Models

    main

    ModelScan operates through a pipeline designed to inspect model files for serialization attacks.

    The Pipeline Lifecycle

    1. Model Discovery: The scanner iterates through the provided path. If it's a directory, it recursively finds files. It also handles zip files by extracting their contents to inspect nested files.
    2. Middleware Execution: For every discovered Model, the MiddlewarePipeline is executed. Middlewares allow for pre-processing or augmenting the model data before scanning.
    3. Scanner Execution: The scanner iterates through all enabled ScanBase implementations. Each scanner performs its specific security checks on the Model object.
    4. Result Aggregation: Issues, errors, and skipped files are collected into internal collections (Issues, errors, skipped) and finally compiled into a structured report.

    Core Abstractions

    • Model: An abstraction representing the file or stream being scanned. It provides access to the model's source and data stream.
    • Scanner (ScanBase): A specialized component that implements the scan(model) method to look for specific vulnerabilities.
    • Middleware: Components that run via a MiddlewarePipeline to intercept or modify the scanning process for a model.
    • Issues: A collection of security findings, categorized by IssueSeverity.

  12. Implement custom issue details with IssueDetails

    main

    If you are building a custom scanner or need to provide additional context for a finding, you must implement the IssueDetails abstract base class. You are required to implement the following methods:

    • output_lines(): Returns a List[str] for human-readable console output.
    • output_json(): Returns a Dict[str, str] for machine-readable JSON output.

    OperatorIssueDetails is a provided implementation specifically for reporting unsafe operator usage.

    from modelscan.issues import IssueDetails
from typing import List, Dict

class MyCustomDetails(IssueDetails):
    def __init__(self, info: str, scanner: str = ""):
        super().__init__(scanner)
        self.info = info

    def output_lines(self) -> List[str]:
        return [f"Custom info: {self.info}"]

    def output_json(self) -> Dict[str, str]:
        return {"info": self.info}

# Usage
details = MyCustomDetails(info="Something suspicious", scanner="my_scanner")