Fickling新型AI/ML pickle文件扫描器:用安全导入白名单守护机器学习供应链

本文详细介绍Trail of Bits团队针对AI/ML领域开发的Fickling pickle文件扫描器新功能。通过创新性的安全导入白名单机制,该工具能有效检测恶意pickle文件,防止远程代码执行攻击。文章包含技术实现原理、验证数据及一行代码集成方案,为机器学习开发者提供实用的供应链安全解决方案。

pickle文件的持续威胁

pickle文件在AI/ML生态系统中仍然是个严重问题,主流机器学习框架的普遍使用不仅增加了模型托管方的远程代码执行(RCE)风险,还使用户面临间接攻击。当用户从Hugging Face等公共平台下载模型时,他们对下载文件中可能包含的恶意文件几乎没有任何防护。

现有工具如Picklescan、ModelScan和model-unpickler可以扫描模型文件并检查危险导入,部分工具甚至直接集成到Hugging Face平台,通过添加安全标签警告用户。但当前扫描器仍存在明显局限,我们通过向Hugging Face测试仓库上传一个未被检测出的恶意pickle文件证实了这一点。

Fickling过滤ML pickle文件的新方法

现有扫描器都依赖检查pickle文件中已知的不安全导入,这种方法本质上受限,因为需要列出所有Python库的可能导入。为突破这一限制,我们团队实现了检测不安全pickle文件的替代方案。

Fickling的新扫描器使用明确的安全导入白名单,仅允许已知安全的导入并阻止其余内容。这一方法基于两项关键研究支撑:

首先,我们确认白名单方法足以过滤所有危险导入并阻断已知pickle利用技术。研究发现,当pickle文件仅包含"安全"导入时无法携带攻击载荷,安全导入必须满足以下所有条件:

  • 不能执行代码或导致代码执行
  • 不能获取或设置对象属性
  • 不能导入其他Python对象或从pickle虚拟机获取已加载对象的引用
  • 不能调用后续反序列化例程

其次,我们确认白名单方法在实践中适用于ML pickle文件。通过分析Hugging Face热门模型的pickle文件,发现它们大多使用相同的少量导入,这意味着可以构建一个小型白名单来覆盖主流公共模型仓库的文件。

我们在3000个Hugging Face顶级仓库的pickle文件上实现了Fickling的ML白名单,通过基准测试验证:Fickling检测出100%的恶意文件,并对99%的安全文件正确分类。

如何使用Fickling新扫描器

我们实现了面向用户的自动pickle验证功能,只需一行代码即可启用。该功能会挂钩pickle模块,使用Fickling的自定义unpickler动态检查加载pickle文件时的每个导入尝试,对未授权导入抛出异常。

使用Fickling保护非常简单,在Python程序开头添加:

1
2
3
import fickling
# 设置全局pickle钩子
fickling.hook.activate_safe_ml_environment()

我们还为用户提供了灵活性:

  • 可在代码库不同位置按需启用/禁用保护
  • 可自定义白名单处理误报情况
  • 我们将持续扩展白名单以减少误报

尽量避免使用pickle

我们的目标是帮助开发者保护系统和AI/ML流水线,欢迎社区对Fickling的AI/ML安全功能提供反馈。如果你正在使用基于pickle的模型,建议尝试该工具。但请记住,避免pickle漏洞的最佳方法是完全避免使用pickle,优先选择基于SafeTensors等更安全格式的模型。

comments powered by Disqus
使用 Hugo 构建
主题 StackJimmy 设计