pickle 是 Python 中用于序列化和反序列化对象的模块。序列化是将对象转换为字节流的过程，反序列化是将字节流还原为对象的过程。它常被用来：

• 存储 Python 对象（变量、列表、字典、类实例等）
• 用来存储模型参数（但不推荐！LLM 现在更常用 safetensors）

import pickle
import os
import torch.nn as nn


class Demo(nn.Module):
    def __init__(self):
        super(Demo, self).__init__()
        self.linear1 = nn.Linear(128, 256)
        self.linear2 = nn.Linear(128, 256)

def test():
    pickle.dump(Demo(), open('demo.bin', 'wb'))
    demo = pickle.load(open('demo.bin', 'rb'))


if __name__ == '__main__':
    test()

但是，使用 Pickle 存在是一个严重的安全风险。在反序列化时，Pickle 可以执行存储在对象中的任意 Python 代码。 这意味着攻击者可以构造恶意的 .pkl 或 .bin 文件，如果你的模型文件是从陌生网站来源下载的，攻击者可以植入恶意代码。

1. 漏洞描述

在 Python 中，pickle 模块用于序列化和反序列化对象，通常使用 pickle.dump 和 pickle.load 进行操作。然而，这个过程中存在一个安全隐患，主要源自 __reduce__ 魔术方法。

接下来，先了解下__reduce__ 魔术方法的作用，它主要是用于在反序列化恢复对象的时候调用（默认不需要手动添加，只有当特殊情况下才会手动编写）。

import pickle

class Demo:
    def __init__(self, a, b):
        self.a = a
        self.b = b

    def __reduce__(self):
        return (self.__class__, (self.a, self.b))

def test():
    # 将 __class__ 对应的类名及其所在模块路径存储到文件中
    # 将 self.a 和 self.b 对应的值存储到文件中
    pickle.dump(Demo(10, 20), open('demo.bin', 'wb'))
    # 从文件中加载数据，并根据存储的模块路径找到名字为 Demo 的类
    # 使用该类和存储的参数值创建对象
    demo = pickle.load(open('demo.bin', 'rb'))

if __name__ == '__main__':
    test()

当一个类实现了 __reduce__ 魔术方法时，pickle 在序列化该对象时会调用 __reduce__，该方法应返回一个元组，其中第一个元素是一个可调用对象，第二个元素是传递给该对象的参数元组。在反序列化时，pickle 会调用这个可调用对象，并传入提供的参数来重建对象。

那么，如果别有用心者在 __reduce__ 返回值上做手脚，即：返回一个具有恶意行为的可调用对象。一旦不知情的用户加载这个 pickle 文件，恶意代码便会执行，从而造成安全威胁。

2. 场景复现

再次明确：Pickle 的安全漏洞主要是由于别有用心者滥用了 __reduce__ 函数，并将返回的正常的可调用对象替换为具有恶意行为的可调用对象。

import pickle
import os
import torch.nn as nn


# 恶意代码的字符串表示
malicious_code_str = """
def malicious_code():
    print("no safe code！")

malicious_code()
"""


class Demo:
    def __reduce__(self):
        return (exec, (malicious_code_str,))
        # return (os.system, ('echo no safe command!',))


def test():
    demo = Demo()
    # 序列化包含恶意程序的对象
    with open('no-safe.pkl', 'wb') as f:
        pickle.dump(demo, f)


if __name__ == '__main__':
    test()

这段代码已经生成了 no-safe.pkl 文件，接下来反序列化该文件：

import pickle

def test():
    demo = pickle.load(open('no-safe.pkl', 'rb'))

if __name__ == '__main__':
    test()

no safe code！

最后总结下，pickle 是一个强大的序列化模块，但在反序列化时存在安全风险。始终确保反序列化的数据来源可靠，避免反序列化用户输入的数据。如果安全性是关键需求，建议使用更安全的序列化格式（如 json、safetensors）。如果必须使用 pickle，可以通过自定义反序列化逻辑来限制可反序列化的类。

import pickle

# 恶意代码的字符串表示
malicious_code_str = """
def malicious_code():
    print("no safe code！")

malicious_code()
"""

class Demo:
    def __reduce__(self):
        return (exec, (malicious_code_str,))
        # return (self.__class__, ())

# 自定义反序列化类
class SafeUnpickler(pickle.Unpickler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 定义允许反序列化的类
        self.allowed_classes = {'__main__.Demo': Demo}

    # 重写 find_class 函数，该函数在反序列化时被调用，用于根据模块名和类名找到对应的类
    def find_class(self, module, name):
        # module 和 name 是从序列化文件中读取的模块和类型
        # 检查模块和类名是否允许反序列化该类
        allowed_class = f'{module}.{name}'
        if allowed_class not in self.allowed_classes:
            raise pickle.UnpicklingError(f'尝试去反序列化不支持的类型 {allowed_class}')
        return self.allowed_classes[allowed_class]


def test():
    demo = Demo()
    pickle.dump(demo, open('no-safe.pkl', 'wb'))
    # 使用自定义反序列化
    with open('no-safe.pkl', 'rb') as f:
        demo = SafeUnpickler(f).load()
    print(demo)

if __name__ == '__main__':
    test()