探索Fickling新特性：强化机器学习系统安全

我们为Fickling添加了新功能，以提供跨广泛机器学习（ML）工作流的增强威胁检测和分析能力。Fickling是一个针对Python pickle模块的反编译器、静态分析器和字节码重写器，可帮助您检测、分析或创建恶意pickle文件。

尽管ML社区已出现更安全的序列化方法（如safetensors文件格式），但pickle的普及所带来的安全风险远未解决。pickle在ML生态系统中的持续广泛采用，使得ML模型文件成为后门、勒索软件、反向shell和其他恶意负载的攻击向量，因此有效识别和缓解此问题至关重要。

为此，我们添加了以下新功能：

• 模块化分析API：生成详细结果，分析pickle文件的恶意行为，并提供方便的JSON输出。
• PyTorch模块：静态分析PyTorch文件并注入代码。
• Polyglot模块：区分、识别和创建不同PyTorch文件格式的多格式文件。

ICYMI：据我们所知，Fickling是第一个针对ML用例定制的pickle安全工具。我们最初的博客文章详细说明了为什么ML pickle文件可被利用，以及Fickling如何具体解决此问题。我们强调，Fickling在潜在恶意文件上运行是安全的，因为它使用自己的Pickle Machine（PM）实现符号执行代码。这使得事件响应和ML基础设施工程师可以使用和部署Fickling，将新颖的ML威胁检测和分析集成到他们的管道中。例如，Fickling已被用于分析在野外发现的恶意ML模型。

模块化分析API

恶意pickle文件可以包含混淆机制以绕过直接扫描。然而，Fickling通过使用其模块化分析API对反编译表示进行静态分析，促进对此类文件的彻底分析。

此API提供详细、系统的方法，将分析分解为特定类别的恶意行为，从而易于确定文件被标记的方式和原因。这使得Fickling成为检查和评估模型工件的有效工具，无论您是想在项目中使用模型之前检查它，还是在妥协后调查工件。

分析封装在易于使用的JSON输出格式中，可从CLI和Python API访问。输出详细说明文件的严重性，提供评估理由，并 pinpoint 触发的特定分析类以及任何相关工件。这种统一的输出格式提高了模块化分析API的可用性，使其易于在不同工具和工作流中定制和集成检测过程。

例如，来自Fickling的Numpy PoC生成的示例恶意pickle文件的输出（图1）：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

{
    "severity": "LIKELY_OVERTLY_MALICIOUS",
    "analysis": "`from posix import system` is suspicious and indicative of an overtly malicious pickle file. Variable `_var0` is assigned value `system(...)` but unused afterward; this is suspicious and indicative of a malicious pickle file",
    "detailed_results": {
        "AnalysisResult": {
            "UnsafeImports": "from posix import system",
            "UnusedVariables": [
                "_var0",
                "system(...)"
            ]
        }
    }
}

图1：Fickling对来自numpy_poc.py的恶意pickle文件分析的JSON输出

PyTorch模块

PyTorch作为最流行的ML框架之一，是ML工作流的重要组成部分。该框架依赖于pickle，这使得Fickling成为携带火炬的绝佳选择。Fickling的PyTorch模块可帮助您处理这些文件。更具体地说，此模块将Fickling的反编译、静态分析和注入功能扩展到PyTorch文件，以便您可以应用模块化分析API和其他功能。这拓宽了Fickling评估生产系统中pickles影响的能力。

在图2中，我们演示了如何使用此PyTorch模块。使用Fickling将保存为PyTorch文件的ML模型转换并序列化为恶意文件。此示例说明了此模块实现的许多用例之一——注入。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

import torch
import torchvision.models as models

from fickling.pytorch import PyTorchModelWrapper

# Load example PyTorch model
model = models.mobilenet_v2()
torch.save(model, "mobilenet.pth")

# Wrap model file into fickling
result = PyTorchModelWrapper("mobilenet.pth")

# Inject payload, overwriting the existing file instead of creating a new one
temp_filename = "temp_filename.pt"
result.inject_payload(
    "print('!!!!!!Never trust a pickle!!!!!!')",
    temp_filename,
    injection="insertion",
    overwrite=True,
)

# Load file with injected payload
# This outputs “!!!!!!Never trust a pickle!!!!!!”. 
torch.load("mobilenet.pth")

图2：Fickling将任意代码注入PyTorch模型文件

Polyglot模块

什么是PyTorch文件？

在深入探讨Polyglot模块之前，让我们再谈谈PyTorch文件。PyTorch文件包含多种不同的文件格式。然而，一个常见的误解是PyTorch文件仅指一种特定的文件格式。格式之间的不当区分会阻碍检测和分析工作，并助长使用这些文件的漏洞利用。Fickling可以区分这些格式，以便在现实世界部署中有效分析它们。

Fickling可以识别以下文件格式：

• PyTorch v0.1.1：包含sys_info、pickle、storages和tensors目录的Tar文件
• PyTorch v0.1.10：堆叠的pickle文件
• TorchScript v1.0：包含model.json文件的ZIP文件
• TorchScript v1.1：包含model.json和attributes.pkl文件（一个pickle文件）的ZIP文件
• TorchScript v1.3：包含data.pkl和constants.pkl文件（两个pickle文件）的ZIP文件
• TorchScript v1.4：包含data.pkl、constants.pkl和版本文件设置为2或更高（两个pickle文件）的ZIP文件
• PyTorch v1.3：包含data.pkl（一个pickle文件）的ZIP文件
• PyTorch模型存档格式[ZIP]：包含Python代码文件和pickle文件的ZIP文件

此列表可能会更改，我们会根据需要不断添加更多文件格式。如果您有兴趣探索PyTorch文件之外的ML文件格式空间，请查看我们全面的ML文件格式列表。

PyTorch文件格式在结构和出现上下文上均有所不同。Polyglot模块的文件格式识别功能可帮助您确保在正确的上下文中使用正确的文件：

• torch.load函数解析PyTorch v1.3、TorchScript v1.4、PyTorch v0.1.10和PyTorch v0.1.1文件。PyTorch v1.3文件格式是这些中最常见的格式，通常被视为规范文件格式。
• 同时，TorchServe系统依赖于PyTorch模型存档格式。
• 已弃用的文件格式（如TorchScript v1.1）被故意包含在Fickling中，因为这些格式仍可能与外部解析器兼容并可能被利用。

图3展示了Fickling如何识别不同的PyTorch文件格式。我们使用torch.save将PyTorch模型序列化为PyTorch v1.3文件和PyTorch v0.1.10文件。Fickling可以清楚地区分这两种不同的格式。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

> import torch
> import torchvision.models as models
> import fickling.polyglot as polyglot
> model = models.mobilenet_v2()
> torch.save(model, "mobilenet.pth")
> polyglot.identify_pytorch_file_format("mobilenet.pth", print_results=True)
Your file is most likely of this format:  PyTorch v1.3 
> torch.save(model, "legacy_mobilenet.pth", _use_new_zipfile_serialization=False)
> polyglot.identify_pytorch_file_format("legacy_mobilenet.pth",
print_results=True)
Your file is most likely of this format:  PyTorch v0.1.10

图3：Fickling区分PyTorch v1.3文件和PyTorch v0.1.10文件

多格式文件？在我的PyTorch中？比您想象的更可能

多格式文件是可以有效解释为多种文件格式的文件。它们已被用于绕过代码签名检查并分发恶意软件，以及其他许多不良行为。您可以在我们关于PolyFile和PolyTracker的博客文章中了解更多关于多格式文件和其他 unruly 解析器的副产品。Fickling对PyTorch文件格式的识别具有多格式感知能力，因为您可以在这些文件之间制作多格式文件。这提出了一个问题：为什么我们应该关心ML模型文件的多格式文件？

多格式ML模型文件可以绕过ML工具中的检查并渗透模型中心，误导该模型的消费者。具体来说，在ML模型文件的上下文中，多格式文件可以是后门ML模型的向量。您可以构建一个多格式文件，使其在解析为一种文件格式时是良性模型，但在解析为另一种文件格式时是后门模型。在我们对safetensors的审计期间，一个现已解决的发现允许我们使用safetensors文件创建多个多格式文件（有趣的事实：报告本身是一个PDF/ZIP多格式文件，ZIP文件包含审计中的多格式文件）。

无论您是在妥协后分析模型工件的多格式性，还是为处理模型文件的MLOps工具构建严格、定义明确的解析器，识别此威胁都很重要。广泛地说，Fickling的Polyglot模块可以帮助我们开始确定多格式文件对ML生态系统的潜在影响。

Fickling还支持创建这些多格式文件用于测试和演示。例如，我们可以使用Fickling制作一个可以有效解释为PyTorch v0.1.10文件和PyTorch模型存档（MAR）文件的文件。

由于pickle是一种流格式，一旦到达STOP操作码就会停止解析，我们可以在不中断有效解析的情况下将任意数据附加到pickle文件。类似地，许多ZIP解析器不强制指定的魔术字节从偏移0开始，这允许我们在保留有效解析的同时将数据预置到ZIP文件。这两种能力结合时，允许我们构建一个既是有效pickle文件又是有效ZIP文件的文件——一个pickle/ZIP多格式文件！

回想一下，PyTorch v0.1.10文件由堆叠的pickles组成。PyTorch MAR解析器是许多接受带有预置数据的文件的ZIP解析器之一。这意味着我们可以基于pickle/ZIP多格式文件，通过将MAR文件附加到PyTorch v0.1.10文件来制作PyTorch v0.1.10 / PyTorch MAR多格式文件。此过程在Fickling中捕获，如以下示例所示：

1
2
3
4

> import fickling.polyglot as polyglot 
> polyglot.create_polyglot("mar_example.mar","legacy_example.pt") 
Making a PyTorch v0.1.10/PyTorch MAR polyglot
The polyglot is contained in polyglot.mar.pt

图4：Fickling创建PyTorch v0.1.10 / PyTorch MAR多格式文件

生成的文件可以使用Fickling准确识别，如下所示：

1
2
3
4

> import fickling.polyglot as polyglot 
> polyglot.identify_pytorch_file_format('polyglot.mar.pt',print_results=True)
Your file is most likely of this format: PyTorch v0.1.10
It is also possible that your file can be validly interpreted as: [‘PyTorch model archive format’]

图5：Fickling识别PyTorch v0.1.10 / PyTorch MAR多格式文件

贡献给Fickling

我们正在积极维护并为Fickling添加新功能，包括新的注入方法、分析类和多格式组合。我们希望Fickling成为进攻和防御安全都可用的工具，因此我们邀请您通过在我们的GitHub上提出问题或直接在我们的联系我们页面上联系我们分享您的反馈。

超越Fickling

虽然Fickling可以帮助您识别由恶意pickle文件引起的对ML系统的威胁，但我们建议完全远离pickle。受限的unpickler可能看起来有用，但它们不是万无一失的解决方案。为了帮助生态系统从pickles向前发展，我们审计了一个更安全的替代方案safetensors；报告了开源代码库中的pickle漏洞；并编写了Semgrep规则以捕获ML库中底层pickling的实例。

我们致力于提高ML生态系统的整体安全性和完整性。请密切关注即将发布的关于保护ML系统的博客文章。