PyTorch 모델은 실행 파일이다: 왜 나는 ‘Pickle Bombs’를 막기 위해 스캐너를 만들었는가

Source: Dev.to

우리 스스로에게 하는 거짓말

DevOps나 보안 분야에 있다면 requirements.txt 를 종교처럼 스캔하고 있을 겁니다. Snyk이나 Dependabot을 사용해 requests 의 취약한 버전을 잡아내고, 안심합니다.

그런 뒤에 데이터 사이언스 팀이 무작위 Hugging Face 저장소에서 5 GB 규모의 .pt (PyTorch) 파일을 다운로드받아 바로 프로덕션 환경에 로드하도록 허용합니다.

우리는 AI 모델을 데이터—무해하고 정적인 파일—처럼 취급합니다.

그렇지 않습니다. 실행 파일입니다.

기술적 현실: .pt 파일이란?

표준 PyTorch 모델 파일은 실제로 ZIP 아카이브입니다. 그 안에 모델 가중치와 구조가 Python의 pickle 모듈을 사용해 직렬화돼 저장됩니다. pickle 모듈은 보안 분야에서 원격 코드 실행 (RCE) 으로 유명합니다.

torch.load('model.pt') 를 실행하면, 언픽클러가 opcode 스트림을 파싱합니다. 그 스트림에 os 를 임포트하고 system() 을 호출하라는 명령이 포함돼 있으면, 머신은 즉시 실행합니다. 샌드박스도, 경고도 없습니다.

이것은 이론이 아닙니다. 5줄의 파이썬 코드로 “악성 모델”을 만드는 방법은 다음과 같습니다:

import torch
import os

class Malicious(object):
    def __reduce__(self):
        return (os.system, ("rm -rf / --no-preserve-root",))

torch.save(Malicious(), "bert_finetune.pt")

그 파일을 로드하면 서버가 사라집니다.

표준 SCA 도구가 실패하는 이유

대부분의 팀은 Dependabot이나 일반 SBOM 생성기 같은 소프트웨어 구성 분석(SCA) 도구에 의존합니다. 이 도구들은 requirements.txt 에 명시된 라이브러리 버전을 알려진 취약점(CVE) 데이터베이스와 매칭합니다.

AI 모델에 이 접근법이 통하지 않는 이유는 다음과 같습니다:

• 라이브러리가 아니다: 커스텀 모델 파일은 버전 번호가 있는 “패키지”가 아니라 바이너리 블롭입니다.
• CVE가 아니다: Pickle Bomb은 공개된 취약점이 아니라 파일 구조 안에 숨겨진 맞춤형 익스플로잇 페이로드입니다.
• 맹점: 의존성에 대한 깔끔한 보안 보고서를 받지만, 5 GB 규모의 RCE 페이로드가 Docker 컨테이너 안에 그대로 존재합니다.

AIsbom 소개

AI 아티팩트의 고유한 특성을 고려한 도구가 필요하다는 것을 깨달았습니다. 우리는 래퍼만 스캔하는 것이 아니라 뇌 자체를 스캔해야 합니다.

AIsbom은 오픈소스 CLI로 깊은 바이너리 내역 분석을 수행합니다:

• 아티팩트를 압축 해제한다: 무거운 가중치를 로드하지 않고 메모리 상에서 PyTorch/Safetensors 구조를 파싱합니다.
• 바이트코드를 역어셈블한다: pickletools 를 사용해 opcode 스트림을 순회합니다.
• 폭탄을 탐지한다: os.system, subprocess, eval, socket 과 같은 위험한 전역 함수를 플래그합니다.

지금 바로 모델을 검사하는 방법

PyPI에서 CLI를 설치합니다:

pip install aisbom-cli

그 다음 모델 폴더를 지정합니다:

aisbom scan ./my-models

다음과 같은 위험도 평가 결과가 출력됩니다:

미래는 SafeTensors (하지만 아직은 아니다)

업계는 .safetensors 라는 안전하고 JSON‑헤더 기반 포맷으로 이동하고 있습니다. 그러나 수백만 개의 레거시 .pt 파일이 여전히 존재하고, 개발자들은 여전히 이를 다운로드합니다. 모두가 마이그레이션할 때까지 방어 장치가 필요합니다.

AI 공급망 보안이 블랙박스가 되어서는 안 된다고 생각해 이 도구를 오픈소스로 공개했습니다.

• GitHub:
• Live Demo: