By Suha S. Hussain

We’ve added new features to Fickling to offer enhanced threat detection and analysis across a broad spectrum of machine learning (ML) workflows. Fickling is a decompiler, static analyzer, and bytecode rewriter for the Python pickle module that can help you detect, analyze, or create malicious pickle files.

While the ML community has seen the rise of safer serialization methods such as the safetensors file format, the security risk posed by the prevalence of pickle is far from resolved. The persistent widespread adoption of pickle in the ML ecosystem allows ML model files to be attack vectors for backdoors, ransomware, reverse shells, and other malicious payloads, making it important that we effectively identify and mitigate this issue.

To that end, we’ve added the following new features:

  1. Modular analysis API: Generate detailed results analyzing pickle files for malicious behaviors, with convenient JSON outputs.
  2. PyTorch module: Statically analyze and inject code into PyTorch files.
  3. Polyglot module: Differentiate, identify, and create polyglots for the different PyTorch file formats.

Pickle overlaying Python code snippet for the fickling tool

ICYMI: To our knowledge, Fickling was the first pickle security tool tailored for ML use cases. Our original blog post detailed why ML pickle files are exploitable and how Fickling specifically addresses this issue. We highlighted that Fickling is safe to run on potentially malicious files because it symbolically executes code using its own implementation of the Pickle Machine (PM). This enables Fickling to be used and deployed by incident response and ML infrastructure engineers to integrate novel ML threat detection and analysis into their pipelines. For instance, Fickling has been used to analyze malicious ML models found in the wild.

Modular analysis API

Malicious pickle files can incorporate obfuscation mechanisms to bypass direct scanning. However, Fickling facilitates a thorough analysis of such files by performing static analysis on the decompiled representations using its modular analysis API.

This API offers a detailed, systematic approach that dissects the analysis into specific categories of malicious behavior so that it’s easy to determine how and why a file was flagged. This makes Fickling an effective tool for inspecting and evaluating model artifacts whether you want to examine a model before using it in a project or investigate artifacts post-compromise.

The analysis is encapsulated in an easy-to-use JSON output format, accessible from both the CLI and Python API. The output details the severity of the file, provides a rationale for its assessment, and pinpoints specific analysis classes that were triggered, along with any relevant artifacts. This unified output format improves the usability of the modular analysis API, making it easy to customize and integrate the detection process across different tools and workflows.

Take, for example, the output from a sample malicious pickle file, generated by Fickling’s Numpy PoC (figure 1):

  • The severity field indicates that Fickling has labeled this file LIKELY_OVERTLY_MALICIOUS.
  • The analysis field explains why: Fickling detected both an unsafe import and an unused variable. The former is a much stronger determinant of severity than the latter. However, not only does detecting the unused variable provide more insight into the use of the unsafe import, but including more granular elements in analysis is especially useful for artifacts that are designed to evade detection.
  • The detailed_results field, expanding upon the analysis field, clearly indicates that the UnsafeImports and UnusedVariables analysis classes were triggered by this file and includes the artifact that triggered both classes. This information can help users make informed decisions based on Fickling’s analysis.
{
    "severity": "LIKELY_OVERTLY_MALICIOUS",
    "analysis": "`from posix import system` is suspicious and indicative of an overtly malicious pickle file. Variable `_var0` is assigned value `system(...)` but unused afterward; this is suspicious and indicative of a malicious pickle file",
    "detailed_results": {
        "AnalysisResult": {
            "UnsafeImports": "from posix import system",
            "UnusedVariables": [
                "_var0",
                "system(...)"
            ]
        }
    }
}

Figure 1: The JSON output of Fickling’s analysis of the malicious pickle file from numpy_poc.py

PyTorch module

PyTorch, one of the most popular frameworks for ML, is an integral component of ML workflows. This framework is dependent on pickle, which makes Fickling an excellent choice for carrying the torch. Fickling’s PyTorch module can help you dill with these files. More concretely, this module extends Fickling’s decompilation, static analysis, and injection capabilities to PyTorch files so you can apply the modular analysis API and other features. This broadens Fickling’s capacity to assess the impact of pickles in production systems.

In figure 2, we demonstrate how this PyTorch module can be used. An ML model saved as a PyTorch file is transformed and serialized into a malicious file using Fickling. This example illustrates just one of the many use cases made possible by this module—injections.

import torch
import torchvision.models as models

from fickling.pytorch import PyTorchModelWrapper

# Load example PyTorch model
model = models.mobilenet_v2()
torch.save(model, "mobilenet.pth")

# Wrap model file into fickling
result = PyTorchModelWrapper("mobilenet.pth")

# Inject payload, overwriting the existing file instead of creating a new one
temp_filename = "temp_filename.pt"
result.inject_payload(
    "print('!!!!!!Never trust a pickle!!!!!!')",
    temp_filename,
    injection="insertion",
    overwrite=True,
)

# Load file with injected payload
# This outputs “!!!!!!Never trust a pickle!!!!!!”. 
torch.load("mobilenet.pth")

Figure 2: Fickling injects arbitrary code into a PyTorch model file.

Polyglot module

What are PyTorch files?

Before we dive into the Polyglot module, let’s talk a bit more about PyTorch files. PyTorch files encompass multiple different file formats. It is a common misconception, however, that a PyTorch file refers to only one specific file format. Improper differentiation between formats hampers detection and analysis efforts and aids exploits that use these files. Fickling can differentiate these formats so that they can be effectively analyzed when used in real-world deployments.

Fickling can identify the following file formats:

  1. PyTorch v0.1.1: Tar file with the sys_info, pickle, storages, and tensors directories
  2. PyTorch v0.1.10: Stacked pickle files
  3. TorchScript v1.0: ZIP file with the model.json file
  4. TorchScript v1.1: ZIP file with the model.json and attributes.pkl files (one pickle file)
  5. TorchScript v1.3: ZIP file with the data.pkl and constants.pkl files (two pickle files)
  6. TorchScript v1.4: ZIP file with the data.pkl, constants.pkl, and version files set at 2 or higher (two pickle files)
  7. PyTorch v1.3: ZIP file containing data.pkl (one pickle file)
  8. PyTorch model archive format [ZIP]: ZIP file that includes Python code files and pickle files

This list is subject to change and we’re continually adding more file formats as needed. If you’re interested in exploring the space of ML file formats beyond PyTorch files, check out our comprehensive list of ML file formats.

The PyTorch file formats differ both in structure and in the contexts where they appear. The Polyglot module’s file format identification feature can help you ensure that the correct files are being used in the correct contexts:

  • The torch.load function parses PyTorch v1.3, TorchScript v1.4, PyTorch v0.1.10, and PyTorch v0.1.1 files. The PyTorch v1.3 file format is the most common format of these and is typically deemed the canonical file format.
  • Meanwhile, TorchServe systems rely on the PyTorch model archive format.
  • Deprecated file formats such as TorchScript v1.1 are deliberately included in Fickling because these formats can still be compatible with external parsers and potentially exploitable.

Figure 3 showcases how Fickling can identify different PyTorch file formats. We used torch.save to serialize a PyTorch model as a PyTorch v1.3 file and a PyTorch v0.1.10 file. Fickling can clearly distinguish these two different formats.

> import torch
> import torchvision.models as models
> import fickling.polyglot as polyglot
> model = models.mobilenet_v2()
> torch.save(model, "mobilenet.pth")
> polyglot.identify_pytorch_file_format("mobilenet.pth", print_results=True)
Your file is most likely of this format:  PyTorch v1.3 
> torch.save(model, "legacy_mobilenet.pth", _use_new_zipfile_serialization=False)
> polyglot.identify_pytorch_file_format("legacy_mobilenet.pth",
print_results=True)
Your file is most likely of this format:  PyTorch v0.1.10

Figure 3: Fickling distinguishes between a PyTorch v1.3 file and a PyTorch v0.1.10 file.

Polyglots? In my PyTorch? It’s more likely than you think

Polyglot files are files that can be validly interpreted as more than one file format. They have been used to bypass code-signing checks and distribute malware, among many other unwanted behaviors. You can learn more about polyglot files and other byproducts of unruly parsers in our blog post on PolyFile and PolyTracker. Fickling’s identification of PyTorch file formats is polyglot-aware because you can make polyglots between these files. This raises the question: Why should we care about polyglots for ML model files?

Polyglot ML model files can bypass checks in ML tools and infiltrate model hubs to mislead consumers of that model. Specifically, in the context of ML model files, polyglot files can be a vector for backdoored ML models. You can construct a polyglot file so that it is a benign model when parsed as one file format but a backdoored model when parsed as another file format. During our audit of safetensors, a now resolved finding allowed us to create multiple polyglots with safetensors files (fun fact: the report itself is a PDF/ZIP polyglot with the ZIP file containing the polyglots from the audit).

It’s important that this threat can be identified whether you’re analyzing a model artifact post-compromise for polyglottery, or building strict, well-defined parsers for MLOps tools that deal with model files. Broadly, Fickling’s Polyglot module can help us begin to determine the potential impact of polyglot files on the ML ecosystem.

Fickling also supports the creation of these polyglot files for testing and demonstration. For instance, we can use Fickling to make a file that can be validly interpreted as both a PyTorch v0.1.10 file and a PyTorch model archive (MAR) file.

Since pickle is a streaming format that stops parsing as soon as it reaches the STOP opcode, we can append arbitrary data to a pickle file without disrupting valid parsing. In a similar vein, many ZIP parsers don’t enforce the specified magic to start at offset 0, which allows us to prepend data to a ZIP file while preserving valid parsing. These two capabilities, when combined, allow us to construct a file that is both a valid pickle file and a valid ZIP file—a pickle/ZIP polyglot!

Recall that a PyTorch v0.1.10 file is composed of stacked pickles. The PyTorch MAR parser is one of many ZIP parsers that accepts files with prepended data. This means that we can build on the pickle/ZIP polyglot to make a PyTorch v0.1.10 / PyTorch MAR polyglot by appending the MAR file to the PyTorch v0.1.10 file. This process is captured in Fickling, as shown in this example:

> import fickling.polyglot as polyglot 
> polyglot.create_polyglot("mar_example.mar","legacy_example.pt") 
Making a PyTorch v0.1.10/PyTorch MAR polyglot
The polyglot is contained in polyglot.mar.pt

Figure 4: Fickling creates a PyTorch v0.1.10 / PyTorch MAR polyglot.

The resulting file can be accurately identified using Fickling, as shown below:

> import fickling.polyglot as polyglot 
> polyglot.identify_pytorch_file_format('polyglot.mar.pt',print_results=True)
Your file is most likely of this format: PyTorch v0.1.10
It is also possible that your file can be validly interpreted as: [‘PyTorch model archive format’]

Figure 5: Fickling identifies a PyTorch v0.1.10 / PyTorch MAR polyglot.

Contribute to Fickling

We are actively maintaining and adding new capabilities to Fickling, including new injection methods, analysis classes, and polyglot combinations. We want Fickling to be a usable tool for both offensive and defensive security, so we invite you to share your feedback by raising an issue on our GitHub or reaching out directly on our Contact us page.

Beyond Fickling

While Fickling can help you identify threats to ML systems caused by malicious pickle files, we recommend moving away from pickle entirely. Restricted unpicklers may seem useful, but they are not a foolproof solution. To help the ecosystem move forward from pickles, we’ve audited a safer alternative, safetensors; reported pickle vulnerabilities in open-source codebases; and written Semgrep rules to catch instances of pickling under the hood in ML libraries.

We’re dedicated to improving the overall security and integrity of the ML ecosystem. Keep an eye out for upcoming blog posts on securing ML systems.