Fickling新型AI/ML pickle文件扫描器——强化机器学习供应链安全

pickle文件持续存在的危险

Python pickle文件本质上不安全，但大多数机器学习模型文件格式仍在使用它们。如果您的代码从外部源加载ML模型，可能会面临安全风险。我们刚刚发布了Fickling（我们的pickle文件扫描器和反编译器）的新改进。Fickling可以轻松集成到AI/ML环境中，捕获可能危害ML模型或托管基础设施的恶意pickle文件。

通过简单的一行代码，Fickling可以在加载pickle文件时强制执行安全导入的白名单，有效阻止隐藏在AI模型中的恶意负载。这满足了AI/ML开发者在pickle文件使用仍然是普遍安全问题的生态系统中对更好供应链安全的需求。

在本博客文章中，我们总结了为AI/ML社区定制Fickling所做的更改，并展示如何集成Fickling的新扫描功能以增强供应链安全性。

pickle文件持续存在的危险

pickle文件在AI/ML生态系统中仍然是一个问题，因为主要ML框架的普遍使用不仅增加了模型主机的远程代码执行（RCE）风险，还使用户面临间接攻击。当用户从Hugging Face平台等公共源下载模型时，他们对下载中可能包含的恶意文件几乎没有保护。

存在诸如Picklescan、ModelScan和model-unpickler等工具来扫描模型文件并检查危险导入。其中一些甚至直接集成到Hugging Face平台中，通过在文件旁边添加小标签来警告用户不安全文件。不幸的是，此措施目前效果不够，因为当前扫描器仍然可以轻松绕过。我们通过向Hugging Face上的测试存储库上传一个未被检测到的恶意pickle文件来确认这一点。

Fickling过滤ML pickle文件的新方法

现有扫描器都依赖于检查pickle文件中是否存在已知硬编码的不安全导入来确定它们是否安全。这种方法本质上有限，因为要真正有效，需要列出几乎所有现有Python库的所有可能导入，这在实际中是不可能的。为了克服这个限制，我们的团队实施了检测不安全pickle文件的替代方法。

Fickling的新扫描器不使用要检查的危险导入列表，而是使用包含可以在pickle文件中安全允许的导入的显式导入白名单。这个想法不是直接检测恶意导入，而是只允许一组已知的安全导入并阻止其余部分。这种方法得到两个关键研究的支持。

首先，我们确认白名单方法足以过滤掉所有危险导入并阻止所有已知的pickle利用技术。我们通过研究现有的pickle安全论文和独立博客文章，并得到我们团队自己的知识和能力的支持来做到这一点。我们发现，当pickle文件仅包含“安全”导入时，它不能携带漏洞利用，这意味着导入的对象必须匹配以下所有标准：

它们不能执行代码或导致代码执行，无论格式如何（编译的代码对象、Python源代码、shell命令、自定义钩子设置等）
它们不能获取或设置对象属性或项目
它们不能导入其他Python对象或从pickle VM内获取对加载的Python对象的引用
它们不能调用后续的反序列化例程（例如，封送或递归调用pickle内的pickle），即使是间接的

其次，我们确认白名单方法可以在实践中为ML pickle文件实施。我们下载并分析了Hugging Face上可用的顶级下载公共模型的pickle文件，并注意到它们中的大多数在其pickle文件中使用相同的几个导入。这意味着可以构建一个小的导入白名单，足以覆盖来自流行公共模型存储库的大多数文件。

我们使用来自顶级Hugging Face存储库的3,000个pickle文件实施了Fickling的ML白名单，检查它们的导入并包含无害的导入。为了验证我们的实施，我们构建了一个基准测试，在两个pickle文件集上运行Fickling：一个干净集包含来自公共Hugging Face存储库的pickle文件，第二个合成数据集包含通过将负载注入第一个集合中的文件获得的恶意pickle文件。Fickling捕获了100%的恶意文件，并正确分类了99%的安全文件。我们当前的实施提供了由手动代码审查支持的导入白名单的强大安全保证，同时仍然保持良好的可用性和非常低的误报率。

如何使用Fickling的新扫描器

在测试和验证Fickling的ML白名单后，我们希望使其易于被最多的人使用。为此，我们实施了一个面向用户的自动pickle验证功能，可以通过一行代码启用。它钩住pickle模块以使用Fickling的自定义unpickler，该unpickler在加载pickle文件时动态检查每个导入。自定义unpickler在尝试进行任何未经白名单授权的导入时引发异常，允许用户捕获潜在不安全的文件并根据需要处理它们。

使用此Fickling保护非常简单。只需在Python程序的最开始运行以下内容：

1
2
3


import fickling
# 这会在pickle上设置全局钩子
fickling.hook.activate_safe_ml_environment()

通过将pickle验证功能打包成一行代码，我们希望促进AI/ML开发者和安全团队系统地采用Fickling。我们的团队也意识到没有一刀切的解决方案，我们还为用户提供了极大的灵活性：

如果需要，您可以在代码库中的不同位置随意启用和禁用保护
如果Fickling因为文件包含未经授权的导入而引发警报，但您确定该文件实际上可以安全加载，您可以轻松自定义白名单一次，使该文件将来通过您的管道

请注意，随着我们继续开发Fickling，我们将继续扩展白名单，从而进一步减少误报的数量。

查看Fickling在GitHub上的文档以获取更多细节！

如果可能，请记住避免使用pickle

我们的努力旨在帮助开发者保护他们的系统和AI/ML管道，我们渴望从社区获得关于Fickling的AI/ML安全功能的反馈。如果您当前使用基于pickle的模型，那么您绝对应该尝试一下——如果您有任何想法，请在Fickling的存储库中打开一个问题。但请记住，避免pickle漏洞利用的最佳方法是完全避免使用pickle，并偏好基于更安全格式（如SafeTensors）的模型。