DEV Community

Labdev
Labdev

Posted on

PyTorch Models Are Executables: Why I Built a Scanner to Stop "Pickle Bombs"

#ai #python #opensource #security

The Lie We Tell Ourselves
If you work in DevOps or Security, you probably scan your requirements.txt religiously. You use Snyk or Dependabot to catch that one vulnerable version of requests. You feel safe.

Then, you turn around and let your Data Science team download a 5GB .pt (PyTorch) file from a random Hugging Face repository and load it directly into your production environment.

We treat AI models like "Data" - inert, harmless static files.

They are not. They are Executables.

The Technical Reality: What is a .pt file?

A standard PyTorch model file is actually a Zip archive. Inside that archive, the model weights and architecture are serialized using Python's pickle module. The pickle module is famous for one thing in the security world: Remote Code Execution (RCE).

When you run torch.load('model.pt'), the unpickler parses a stream of opcodes. If that stream contains instructions to import os and run system(), your machine executes it. Instantly. No sandbox. No warning.

This isn't theoretical. Here is how easy it is to create a "Malicious Model" in 5 lines of Python:

import torch
import os

class Malicious(object):
    def __reduce__(self):
        return (os.system, ("rm -rf / --no-preserve-root",))

torch.save(Malicious(), "bert_finetune.pt")
Enter fullscreen mode Exit fullscreen mode

If you load that file, your server is gone.

Why Standard SCA Tools Fail
Most teams rely on Software Composition Analysis (SCA) tools like Dependabot or standard SBOM generators. These tools work by matching library versions in requirements.txt against databases of known vulnerabilities (CVEs).

This approach fails for AI models because:

  • Not a Library: A custom model file isn't a "package" with a version number. It's a binary blob.
  • Not a CVE: A Pickle Bomb isn't a public vulnerability; it's a custom exploit payload hidden inside the file structure.
  • Blind Spot: You get a clean security report for your dependencies, while a 5GB RCE payload sits undetected in your Docker container.

Introducing AIsbom

I realized we needed a tool that respects the unique nature of AI artifacts. We don't just need to scan the wrapper; we need to scan the brain.

So, I built AIsbom.

It is an open-source CLI that performs Deep Binary Introspection.

  • It Unzips the Artifact: It parses PyTorch/Safetensors structures in memory (without loading the heavy weights).
  • It Disassembles the Bytecode: It uses pickletools to iterate over the opcode stream.
  • It Detects the Bomb: It flags dangerous globals like os.system, subprocess, eval, and socket.

How to check your models right now

You can install the CLI from PyPI:

pip install aisbom-cli

Then, point it at your model folder:

aisbom scan ./my-models

You'll get a risk assessment like this:

Filename           | Framework | Risk Level
-------------------|-----------|-------------------------------------

bert_finetune.pt   | PyTorch   | CRITICAL (RCE Detected: posix.system)
safe_model.st      | SafeTensors | LOW (Binary Safe)
Enter fullscreen mode Exit fullscreen mode

The Future is SafeTensors (But we aren't there yet)

The industry is moving toward .safetensors, which is a safe, JSON-header-based format. But millions of legacy .pt files still exist, and developers still download them.

Until everyone migrates, we need guardrails.

I open-sourced this tool because I believe AI Supply Chain Security shouldn't be a black box.

GitHub: https://github.com/Lab700xOrg/aisbom
Live Demo: https://aisbom.io

Let me know what you think. I'm actively looking for edge cases in Pickle protocols to make the detection even more robust.

Top comments (0)