6 ms·

AIsbom – open-source CLI to detect "Pickle Bombs" in PyTorch models

52 points by lab700xdev 9mo ago

eyeris 9mo ago
Was there any research into prior art? Recently did some research into this space and it seems like there are already a bunch of off the shelf open source projects for address this
roywiggins 9mo ago
Don't love how ChatGPT the readme is, the bullet points under "Why AIsbom?" are very, very ChatGPT.
- xpe 9mo ago
  I will preemptively grant the narrow point that if a project demonstrates poor quality in its code or text (i.e. what I mean when I say "slop"), it can dissuade potential users. However, the "Why AIsbom?" section strikes me as clear and informative. Many people prefer human writing. I get it, and I think I understand most of the underlying reasons and emotional drives. [1] Nevertheless, my top preference (I think?) is clarity and accuracy. For technical writing, if these two qualities are present, I'm rarely bothered by what people may label "AI writing". OTOH, when I see sloppy, poorly reasoned, out-of-date writing, my left hand readies itself for ⌘W. [2] A suggestion for the comment above, which makes a stylistic complaint: be more specific about what can be improved. Finally, a claim: over time, valid identification of some text as being AI-generated will require more computation and be less accurate. [3] [1]: Food for thought: https://theconversation.com/people-say-they-prefer-stories-written-by-humans-over-ai-generated-works-yet-new-study-suggests-thats-not-quite-true-251347 https://theconversation.com/people-say-they-prefer-stories-w... and the backing report: https://docs.iza.org/dp17646.pdf https://docs.iza.org/dp17646.pdf [2]: To be open, I might just have a much higher than average bar for precision -- I tend to prefer reading source materials than derivative press coverage, and I prefer reading a carefully worded, dry written documentation file over an informal chat description. To keep digging the hole for myself, I usually don't like the modern practice of putting unrelated full-width pictures in blog posts because they look purdy. Maybe it comes from a "just the facts, please" mentality when reading technical material. [3]: I realize this isn't the clearest testable prediction, but I think the gist of it is falsifiable.
lab700xdev 9mo ago
Hi HN, I’ve been working with ML infrastructure for a while and realized there’s a gap in the security posture: we scan our requirements.txt for vulnerabilities, but blindly trust the 5GB binary model files (.pt) we download from Hugging Face. Most developers don't realize that standard PyTorch files are just Zip archives containing Python Pickle bytecode. When you run torch.load(), the unpickler executes that bytecode. This allows for arbitrary code execution (RCE) inside the model file itself - what security researchers call a "Pickle Bomb." I built AIsbom (AI Software Bill of Materials) to solve this without needing a full sandbox. How it works: 1. It inspects the binary structure of artifacts (PyTorch, Pickle, Safetensors) without loading weights into RAM. 2. For PyTorch/Pickles, it uses static analysis (via pickletools) to disassemble the opcode stream. 3. It looks for GLOBAL or STACK_GLOBAL instructions referencing dangerous modules like os.system, subprocess, or socket. 4. It outputs a CycloneDX v1.6 JSON SBOM compatible with enterprise tools like Dependency-Track. 5. It also parses .safetensors headers to flag "Non-Commercial" (CC-BY-NC) licenses, which often slip into production undetected. It’s open source (Apache 2.0) and written in Python/Typer. Repo: https://github.com/Lab700xOrg/aisbom https://github.com/Lab700xOrg/aisbom Live Demo (Web Viewer): https://aisbom.io https://aisbom.io Why I built a scanner? https://dev.to/labdev_c81554ba3d4ae28317/pytorch-models-are-executables-why-i-built-a-scanner-to-stop-pickle-bombs-92i https://dev.to/labdev_c81554ba3d4ae28317/pytorch-models-are-... I’d love feedback on the detection logic (specifically safety.py) or if anyone has edge cases of weird Pickle protocols that break the disassembler.
- rafram 9mo ago
  > It looks for GLOBAL or STACK_GLOBAL instructions referencing dangerous modules like os.system, subprocess, or socket. This seems like a doomed approach. You can’t make a list of every “dangerous” function in every library.
  - oofbey 9mo ago
    Agree an explicit block list is not very robust. I imagine the vast majority of legit ML models use only a very limited set of math functions and basically no system interaction. Would be good to fingerprint a big set of assumed-safe models and flag anything which diverges from that.
chuckadams 9mo ago
When dealing with stuff like php serialization and pickle, the rule is simple: never unpickle anything you didn't pickle yourself. If anything else could possibly touch the serialized bytes, sign it with HMAC and keep that somewhere untouchable. I somehow doubt this tool is going to be able to pull off what Java bytecode verification could not.
- lab700xdev 9mo ago
  The Golden Rule holds: "Don't unpickle untrusted data." The problem I'm trying to solve is that "Untrusted" has become blurry in the AI age. Data Scientists treat Model Hubs (like Hugging Face) as trusted repositories, similar to PyPI or NPM. They shouldn't, but they do. This tool effectively serves as a "Loud Warning Label" to break that assumption. It tells the engineer: "Hey, you think this is just weights, but I see socket calls in here. Do not load this."
- nextaccountic 9mo ago
  > When dealing with stuff like php serialization and pickle, the rule is simple: never unpickle anything you didn't pickle yourself. I thought the rule was, never use pickle, it makes no sense when other serialization formats exist and are just as easy to use
  - deleted 9mo ago
    [deleted]
woodruffw 9mo ago
The checks here seem pretty minimal[1]. I'd recommend taking a look at fickling (FD: former employer) for a more general approach to pickle decompilation/analysis[2]. [1]: https://github.com/Lab700xOrg/aisbom/blob/main/aisbom/safety.py https://github.com/Lab700xOrg/aisbom/blob/main/aisbom/safety... [2]: https://github.com/trailofbits/fickling https://github.com/trailofbits/fickling
- lab700xdev 9mo ago
  Thanks for the link! fickling is excellent work (and definitely the gold standard for deep analysis). The goal with AIsbom was to build something lightweight enough to run in a fast CI/CD loop that creates a standard inventory (CycloneDX SBOM) alongside the security check. We are definitely looking at fickling's symbolic execution approach for inspiration on how to make our safety.py module more robust against obfuscation.
  - liuliu 9mo ago
    I know this sounds weird: "symbolic execution" of pickle VM cannot be slow right? We are talking about just a few thousands instructions here and you don't need "symbolic execution" per se, just write a custom interpreter and run it. That would take less than 10ms for any given PyTorch file (excluding disk loading).
- liuliu 9mo ago
  Agree. Writing a pickle interpreter is not particularly challenging. I did that in Swift to help load PyTorch checkpoint https://github.com/liuliu/swift-fickling https://github.com/liuliu/swift-fickling without these pitfalls.
deleted 9mo ago
[deleted]
esafak 9mo ago
Good job. Pickle has no place in production. Yeah, I said it.
krick 9mo ago
Ok, as others noted, the tool in question is hardly a solution, but what is, then? I mean, presented like that, it's pretty crazy that everyone just downloads and runs 5GB executable blobs from Hugging Face. Does anyone review them somehow before they are accepted and gather 10K downloads on HF? Or is it really another totally mindbogglingy crazy thing that happens right now across all world, and everybody just shrugs and waits for catastrophic breach of security of planetary scale to happen?

52 points by lab700xdev 9mo ago

Was there any research into prior art? Recently did some research into this space and it seems like there are already a bunch of off the shelf open source projects for address this

roywiggins 9mo ago

Don't love how ChatGPT the readme is, the bullet points under "Why AIsbom?" are very, very ChatGPT.

xpe 9mo ago
I will preemptively grant the narrow point that if a project demonstrates poor quality in its code or text (i.e. what I mean when I say "slop"), it can dissuade potential users. However, the "Why AIsbom?" section strikes me as clear and informative. Many people prefer human writing. I get it, and I think I understand most of the underlying reasons and emotional drives. [1] Nevertheless, my top preference (I think?) is clarity and accuracy. For technical writing, if these two qualities are present, I'm rarely bothered by what people may label "AI writing". OTOH, when I see sloppy, poorly reasoned, out-of-date writing, my left hand readies itself for ⌘W. [2] A suggestion for the comment above, which makes a stylistic complaint: be more specific about what can be improved. Finally, a claim: over time, valid identification of some text as being AI-generated will require more computation and be less accurate. [3] [1]: Food for thought: https://theconversation.com/people-say-they-prefer-stories-written-by-humans-over-ai-generated-works-yet-new-study-suggests-thats-not-quite-true-251347 https://theconversation.com/people-say-they-prefer-stories-w... and the backing report: https://docs.iza.org/dp17646.pdf https://docs.iza.org/dp17646.pdf [2]: To be open, I might just have a much higher than average bar for precision -- I tend to prefer reading source materials than derivative press coverage, and I prefer reading a carefully worded, dry written documentation file over an informal chat description. To keep digging the hole for myself, I usually don't like the modern practice of putting unrelated full-width pictures in blog posts because they look purdy. Maybe it comes from a "just the facts, please" mentality when reading technical material. [3]: I realize this isn't the clearest testable prediction, but I think the gist of it is falsifiable.

lab700xdev 9mo ago

Hi HN, I’ve been working with ML infrastructure for a while and realized there’s a gap in the security posture: we scan our requirements.txt for vulnerabilities, but blindly trust the 5GB binary model files (.pt) we download from Hugging Face. Most developers don't realize that standard PyTorch files are just Zip archives containing Python Pickle bytecode. When you run torch.load(), the unpickler executes that bytecode. This allows for arbitrary code execution (RCE) inside the model file itself - what security researchers call a "Pickle Bomb." I built AIsbom (AI Software Bill of Materials) to solve this without needing a full sandbox. How it works: 1. It inspects the binary structure of artifacts (PyTorch, Pickle, Safetensors) without loading weights into RAM. 2. For PyTorch/Pickles, it uses static analysis (via pickletools) to disassemble the opcode stream. 3. It looks for GLOBAL or STACK_GLOBAL instructions referencing dangerous modules like os.system, subprocess, or socket. 4. It outputs a CycloneDX v1.6 JSON SBOM compatible with enterprise tools like Dependency-Track. 5. It also parses .safetensors headers to flag "Non-Commercial" (CC-BY-NC) licenses, which often slip into production undetected. It’s open source (Apache 2.0) and written in Python/Typer. Repo: https://github.com/Lab700xOrg/aisbom https://github.com/Lab700xOrg/aisbom Live Demo (Web Viewer): https://aisbom.io https://aisbom.io Why I built a scanner? https://dev.to/labdev_c81554ba3d4ae28317/pytorch-models-are-executables-why-i-built-a-scanner-to-stop-pickle-bombs-92i https://dev.to/labdev_c81554ba3d4ae28317/pytorch-models-are-... I’d love feedback on the detection logic (specifically safety.py) or if anyone has edge cases of weird Pickle protocols that break the disassembler.

rafram 9mo ago
> It looks for GLOBAL or STACK_GLOBAL instructions referencing dangerous modules like os.system, subprocess, or socket. This seems like a doomed approach. You can’t make a list of every “dangerous” function in every library.
- oofbey 9mo ago
  Agree an explicit block list is not very robust. I imagine the vast majority of legit ML models use only a very limited set of math functions and basically no system interaction. Would be good to fingerprint a big set of assumed-safe models and flag anything which diverges from that.

chuckadams 9mo ago

When dealing with stuff like php serialization and pickle, the rule is simple: never unpickle anything you didn't pickle yourself. If anything else could possibly touch the serialized bytes, sign it with HMAC and keep that somewhere untouchable. I somehow doubt this tool is going to be able to pull off what Java bytecode verification could not.

lab700xdev 9mo ago
The Golden Rule holds: "Don't unpickle untrusted data." The problem I'm trying to solve is that "Untrusted" has become blurry in the AI age. Data Scientists treat Model Hubs (like Hugging Face) as trusted repositories, similar to PyPI or NPM. They shouldn't, but they do. This tool effectively serves as a "Loud Warning Label" to break that assumption. It tells the engineer: "Hey, you think this is just weights, but I see socket calls in here. Do not load this."
nextaccountic 9mo ago
> When dealing with stuff like php serialization and pickle, the rule is simple: never unpickle anything you didn't pickle yourself. I thought the rule was, never use pickle, it makes no sense when other serialization formats exist and are just as easy to use
- deleted 9mo ago
  [deleted]

woodruffw 9mo ago

The checks here seem pretty minimal[1]. I'd recommend taking a look at fickling (FD: former employer) for a more general approach to pickle decompilation/analysis[2]. [1]: https://github.com/Lab700xOrg/aisbom/blob/main/aisbom/safety.py https://github.com/Lab700xOrg/aisbom/blob/main/aisbom/safety... [2]: https://github.com/trailofbits/fickling https://github.com/trailofbits/fickling

lab700xdev 9mo ago
Thanks for the link! fickling is excellent work (and definitely the gold standard for deep analysis). The goal with AIsbom was to build something lightweight enough to run in a fast CI/CD loop that creates a standard inventory (CycloneDX SBOM) alongside the security check. We are definitely looking at fickling's symbolic execution approach for inspiration on how to make our safety.py module more robust against obfuscation.
- liuliu 9mo ago
  I know this sounds weird: "symbolic execution" of pickle VM cannot be slow right? We are talking about just a few thousands instructions here and you don't need "symbolic execution" per se, just write a custom interpreter and run it. That would take less than 10ms for any given PyTorch file (excluding disk loading).
liuliu 9mo ago
Agree. Writing a pickle interpreter is not particularly challenging. I did that in Swift to help load PyTorch checkpoint https://github.com/liuliu/swift-fickling https://github.com/liuliu/swift-fickling without these pitfalls.

deleted 9mo ago

[deleted]

esafak 9mo ago

Good job. Pickle has no place in production. Yeah, I said it.

krick 9mo ago

Ok, as others noted, the tool in question is hardly a solution, but what is, then? I mean, presented like that, it's pretty crazy that everyone just downloads and runs 5GB executable blobs from Hugging Face. Does anyone review them somehow before they are accepted and gather 10K downloads on HF? Or is it really another totally mindbogglingy crazy thing that happens right now across all world, and everybody just shrugs and waits for catastrophic breach of security of planetary scale to happen?