Draichi/pytorch-pickle-rce-poc

🎓 AI Security Research: PyTorch Pickle Deserialization PoC (harmless payload) 🎓

Purpose of this repository

This is an educational security research artifact, not a functional model. It was built as part of hands-on AI Red Teaming coursework (Hack The Box Academy) to demonstrate, end-to-end, a well-known and already publicly documented class of vulnerability: arbitrary code execution via unsafe pickle deserialization of ML checkpoints. Hugging Face's own security docs describe this exact mechanism — see "Pickle Scanning" — and this repo is a hands-on, disclosed reproduction of that documented risk, built for learning and awareness rather than for use against any target.

Video demo: unedited screen recording of loading this exact checkpoint and watching the payload fire.

To keep the demonstration unambiguous and safe to publish, the embedded "payload" carries zero harmful capability by design: it opens the default web browser to a YouTube video (Rick Astley's "Never Gonna Give You Up"). No network callback, no filesystem access, no credential access, no persistence, no shell. The vulnerability class demonstrated is real and serious; the payload used to illustrate it deliberately is not.

Who this is for: ML engineers, MLOps/platform teams, and security practitioners who want to see — concretely and safely — what "just load the checkpoint" can mean on an untrusted file, so they can make informed decisions about weights_only, safetensors, and model provenance in their own pipelines.

Does torch.load() trigger this automatically?

Not by default, and that matters. Since PyTorch 2.6, torch.load() defaults to weights_only=True, which uses a restricted unpickler and blocks exactly this class of attack. Loading this file with a modern, unmodified torch.load(path) will raise an UnpicklingError and refuse to execute the payload — verified against PyTorch 2.11.

The payload does fire if the checkpoint is loaded via:

• torch.load(path, weights_only=False) — often used to "get an old checkpoint to load" or because the loading code predates PyTorch 2.6.
• Plain pickle.load(open(path, "rb")) — has never had this restriction and is still fully unsafe.
• Any older PyTorch version (< 2.6) using default settings.

In other words: the platform-level fix exists and works, but it's opt-out, not universal. A huge amount of real-world code still explicitly sets weights_only=False, or uses pickle directly, or predates the change.

What this model actually is

Underneath the trap, rickroll_trojan_model.pth wraps a small, legitimately trained SimpleNet (a couple of nn.Linear layers). It has no real-world use as a model — its only purpose is to be loaded.

How the attack works

1. Steganographic payload hiding. A short Python snippet (webbrowser.open(url)) is embedded into the least significant bits (LSBs) of one of the model's weight tensors, 1–2 bits per float32 value. The perturbation is small enough to be statistically invisible — it does not show up as an anomalous file size, hash mismatch, or obvious weight distribution shift.
2. Pickle __reduce__ hijack. The object actually saved to disk is not a plain state_dict, it's a wrapper class whose __reduce__ method returns (exec, (loader_code,)). Python's pickle module calls this automatically during unsafe deserialization, which means the loader code runs as a side effect of unpickling, not of any model call.
3. Trigger. The loader code reconstructs the real state_dict, extracts the poisoned tensor, decodes the hidden payload from its LSBs, and executes it via exec().
4. Result. torch.load("rickroll_trojan_model.pth", weights_only=False) (or plain pickle.load) is sufficient to trigger execution. No model.eval(), no inference call, no user click beyond loading the file.

This is the exact same mechanism used by real-world malicious model attacks (e.g., reverse shells, credential stealers, cryptominers hidden in .pth/.pkl files uploaded to public model hubs) — the only thing swapped out here is the payload.

How to inspect this file safely (i.e., without running it)

Never load an untrusted checkpoint with weights_only=False or plain pickle.load(). To verify what a .pth/.pkl file actually does first:

• Static opcode inspection: use fickling or picklescan to disassemble the pickle stream and flag dangerous opcodes (GLOBAL, REDUCE, calls to exec, os.system, subprocess, etc.) without executing anything.
• Prefer safe formats: safetensors stores only tensor data, no executable pickle opcodes, and cannot do this by construction. If you don't control the source of a checkpoint, insist on .safetensors.
• Leave weights_only at its default (True) on PyTorch >= 2.6. This file is a good test case: loading it this way should fail with an UnpicklingError rather than execute anything.
• Sandbox first: if you must load an unfamiliar checkpoint with weights_only=False from a trusted-but-unverified source, do it in an isolated VM/container with no network access and no sensitive credentials mounted.

Why this matters

Public model hubs (including this one) work on trust: anyone can upload a .pth file, and depending on how it's loaded, that can be equivalent to running arbitrary code on your machine. weights_only=True closes the default path, but plenty of real pipelines still opt out of it, and pickle.load() was never protected. This model is a controlled, disclosed, harmless demonstration of that gap, so that teams can see the failure mode before they encounter it with a payload that isn't a YouTube link.

Intent & responsible use

This checkpoint is published openly and with full disclosure as part of AI Red Teaming study material. The embedded payload is intentionally non-destructive, makes no network connections, exfiltrates nothing, and leaves no persistence. It is safe to trigger deliberately (e.g., in a VM or throwaway environment, explicitly passing weights_only=False) to see the technique in action, but it should not be loaded that way inside any pipeline, notebook, or environment holding real credentials or sensitive data — the point is exactly that you cannot tell from the file itself that this is happening.

Do not strip this model card, rename the file to obscure its origin, or re-upload the payload with anything other than a harmless action. Repurposing this technique against a target without authorization is illegal in most jurisdictions (unauthorized access / computer misuse statutes) — this artifact is for authorized security research, red teaming, and education only.

Repro / source

Built with PyTorch, LSB steganography over float32 weight tensors, and a pickle.__reduce__ loader. Full notebook walkthrough (training the base model, encoding the payload, building the malicious wrapper) is available here:

https://github.com/Draichi/htb/blob/main/data_attacks/steganography_rickroll_attack.ipynb