Pickle 扫描

Pickle 是机器学习中广泛使用的序列化格式。最值得注意的是，它是 PyTorch 模型权重的默认格式。

加载 pickle 文件时可能会发生危险的任意代码执行攻击。我们建议加载来自您信任的用户和组织的模型，依赖于签名的提交，和/或使用 from_tf=True 自动转换机制加载 TF 或 Jax 格式的模型。我们还通过直接在 Hub 上显示/“审查”任何 pickle 文件中的导入列表来缓解此问题。最后，我们正在尝试一种新的、简单的权重序列化格式，称为 safetensors。

什么是 Pickle？

来自官方文档

pickle 模块实现了用于序列化和反序列化 Python 对象结构的二进制协议。

这意味着 pickle 是一种序列化协议，用于在各方之间高效地共享数据。

我们将 pickle 称为在序列化过程中生成的二进制文件。

从本质上讲，pickle 基本上是一堆指令或操作码。正如您可能猜到的，它不是人类可读的。操作码在序列化时生成，并在反序列化时顺序读取。根据操作码，执行给定的操作。

这是一个小例子

import pickle
import pickletools

var = "data I want to share with a friend"

# store the pickle data in a file named 'payload.pkl'
with open('payload.pkl', 'wb') as f:
    pickle.dump(var, f)

# disassemble the pickle
# and print the instructions to the command line
with open('payload.pkl', 'rb') as f:
    pickletools.dis(f)

运行此代码时，它将创建一个 pickle 文件并在您的终端中打印以下指令

    0: \x80 PROTO      4
    2: \x95 FRAME      48
   11: \x8c SHORT_BINUNICODE 'data I want to share with a friend'
   57: \x94 MEMOIZE    (as 0)
   58: .    STOP
highest protocol among opcodes = 4

现在不必太担心这些指令，只需要知道 pickletools 模块对于分析 pickle 非常有用。它允许您读取文件中的指令，**无需**执行任何代码。

Pickle 不仅仅是一种序列化协议，它还通过允许用户在反序列化时运行 Python 代码来提供更大的灵活性。这听起来不太好，对吧？

为什么它很危险？

如上所述，反序列化 pickle 意味着可以执行代码。但这有一些限制：您只能引用顶级模块中的函数和类；您不能将它们嵌入到 pickle 文件本身中。

回到绘图板

import pickle
import pickletools

class Data:
    def __init__(self, important_stuff: str):
        self.important_stuff = important_stuff

d = Data("42")

with open('payload.pkl', 'wb') as f:
    pickle.dump(d, f)

运行此脚本时，我们将再次获得 payload.pkl。当我们检查文件内容时


# cat payload.pkl
__main__Data)}important_stuff42sb.%

# hexyl payload.pkl
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 80 04 95 33 00 00 00 00 ┊ 00 00 00 8c 08 5f 5f 6d │×•×30000┊000×•__m│
│00000010│ 61 69 6e 5f 5f 94 8c 04 ┊ 44 61 74 61 94 93 94 29 │ain__××•┊Data×××)│
│00000020│ 81 94 7d 94 8c 0f 69 6d ┊ 70 6f 72 74 61 6e 74 5f │××}××•im┊portant_│
│00000030│ 73 74 75 66 66 94 8c 02 ┊ 34 32 94 73 62 2e       │stuff××•┊42×sb.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

我们可以看到其中没有太多内容，一些操作码和相关数据。您可能在想，那么 pickle 的问题是什么？

让我们尝试其他方法

from fickling.pickle import Pickled
import pickle

# Create a malicious pickle
data = "my friend needs to know this"

pickle_bin = pickle.dumps(data)

p = Pickled.load(pickle_bin)

p.insert_python_exec('print("you\'ve been pwned !")')

with open('payload.pkl', 'wb') as f:
    p.dump(f)

# innocently unpickle and get your friend's data
with open('payload.pkl', 'rb') as f:
    data = pickle.load(f)
    print(data)

在这里，为了简单起见，我们使用 fickling 库。它允许我们添加 pickle 指令以通过 exec 函数执行字符串中包含的代码。这就是您规避无法在 pickle 中定义函数或类的事实的方法：您对保存为字符串的 Python 代码运行 exec。

运行此代码时，它将创建一个 payload.pkl 并打印以下内容

you've been pwned !
my friend needs to know this

如果我们检查 pickle 文件的内容，我们会得到

# cat payload.pkl
c__builtin__
exec
(Vprint("you've been pwned !")
tR my friend needs to know this.%

# hexyl payload.pkl
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 63 5f 5f 62 75 69 6c 74 ┊ 69 6e 5f 5f 0a 65 78 65 │c__built┊in___exe│
│00000010│ 63 0a 28 56 70 72 69 6e ┊ 74 28 22 79 6f 75 27 76 │c_(Vprin┊t("you'v│
│00000020│ 65 20 62 65 65 6e 20 70 ┊ 77 6e 65 64 20 21 22 29 │e been p┊wned !")│
│00000030│ 0a 74 52 80 04 95 20 00 ┊ 00 00 00 00 00 00 8c 1c │_tR×•× 0┊000000×•│
│00000040│ 6d 79 20 66 72 69 65 6e ┊ 64 20 6e 65 65 64 73 20 │my frien┊d needs │
│00000050│ 74 6f 20 6b 6e 6f 77 20 ┊ 74 68 69 73 94 2e       │to know ┊this×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

基本上，这就是您反序列化时发生的情况

# ...
opcodes_stack = [exec_func, "malicious argument", "REDUCE"]
opcode = stack.pop()
if opcode == "REDUCE":
    arg = opcodes_stack.pop()
    callable = opcodes_stack.pop()
    opcodes_stack.append(callable(arg))
# ...

构成威胁的指令是 STACK_GLOBAL、GLOBAL 和 REDUCE。

REDUCE 会告诉反序列化器使用提供的参数执行函数，而 *GLOBAL 指令会告诉反序列化器 导入内容。

总之，pickle 很危险，因为

导入 Python 模块时，可以执行任意代码
您可以导入内置函数，如 eval 或 exec，这些函数可用于执行任意代码
实例化对象时，可能会调用构造函数

这就是为什么在大多数使用 pickle 的文档中都说明，不要从不受信任的来源反序列化数据。

缓解策略

不要使用 Pickle

Luc 的建议很好，但 pickle 被广泛使用，而且短期内不会消失：找到每个人都满意的新的格式并启动更改需要一些时间。

那么我们现在能做什么？

加载来自您信任的用户和组织的文件

在 Hub 上，您可以使用 GPG 密钥签名您的提交。这并**不能**保证您的文件是安全的，但它确实保证了文件的来源。

如果您认识并信任用户 A，并且包含 Hub 上文件的提交已由用户 A 的 GPG 密钥签名，那么可以安全地假设您可以信任该文件。

加载 TF 或 Flax 模型权重

TensorFlow 和 Flax 检查点不受影响，并且可以使用 from_pretrained 方法的 from_tf 和 from_flax 关键字参数在 PyTorch 架构中加载，以规避此问题。

例如：

from transformers import AutoModel

model = AutoModel.from_pretrained("google-bert/bert-base-cased", from_flax=True)

使用您自己的序列化格式

最后一种格式，safetensors，是我们目前正在开发和试验的一种简单的序列化格式！如果您能提供帮助或贡献，我们将不胜感激 🔥。

改进 torch.load/save

PyTorch 上正在进行一项公开讨论，讨论是否要默认情况下提供一种安全的方法来仅从 *.pt 文件加载权重——请参与讨论！

Hub 的安全扫描器

我们现在拥有的

我们创建了一个安全扫描器，它会扫描推送到 Hub 的每个文件并运行安全检查。在撰写本文时，它运行两种类型的扫描

ClamAV 扫描
Pickle 导入扫描

对于 ClamAV 扫描，文件会通过开源杀毒软件 ClamAV 运行。虽然这涵盖了大量危险文件，但它不涵盖 pickle 漏洞。

我们实现了一个 Pickle 导入扫描，它会提取 pickle 文件中引用的导入列表。每次上传 pytorch_model.bin 或任何其他 pickle 文件时，都会运行此扫描。

在 Hub 上，导入列表将显示在包含导入的每个文件旁边。如果任何导入看起来可疑，则会突出显示。

我们借助 pickletools.genops 获取此数据，它允许我们在不执行潜在危险代码的情况下读取文件。

请注意，这使我们能够知道，在解封文件时，它是否会对通过 *GLOBAL 导入的潜在危险函数进行 REDUCE 操作。

免责声明：这并非万无一失。作为用户，您有责任检查某事物是否安全。我们没有积极地审核 Python 包的安全性，我们拥有的安全/不安全导入列表以尽力而为的方式维护。如果您认为某事物不安全，并且我们将其标记为不安全，请通过发送电子邮件至 huggingface.co 网站联系我们。

潜在解决方案

人们可能会想到像这样创建一个自定义的 Unpickler。但正如我们在这个复杂的漏洞利用中看到的，这行不通。

值得庆幸的是，始终存在 eval 导入的痕迹，因此直接读取操作码应该可以捕获恶意使用。

我目前提出的解决方案是创建一个类似于 .gitignore 但用于导入的文件。

此文件将是一个导入的白名单，如果存在未包含在白名单中的导入，则会将 pytorch_model.bin 文件标记为危险文件。

可以想象使用一种类似正则表达式的格式，例如，您可以通过一行简单的代码允许所有 numpy 子模块：numpy.*。

进一步阅读

pickle - Python 对象序列化 - Python 3.10.6 文档

危险的 Pickles - 恶意的 Python 序列化

GitHub - trailofbits/fickling: Python pickling 反编译器和静态分析器

利用 Python pickles

cpython/pickletools.py at 3.10 · python/cpython

cpython/pickle.py at 3.10 · python/cpython

CrypTen/serial.py at main · facebookresearch/CrypTen

CTFtime.org / Balsn CTF 2019 / pyshv1 / Writeup

恢复 Python 的 pickle 模块

< > 更新在 GitHub 上