SimpleInjector.cpp

// SimpleInjector.cpp : Defines the entry point for the application.
//

#include "framework.h"
#include "SimpleInjector.h"
#include <d3d11.h>
#include <vector>
#include <string>
#include <set>
#include <TlHelp32.h>
#include <commdlg.h>
#include "ImGui/imgui.h"
#include "ImGui/imgui_impl_win32.h"
#include "ImGui/imgui_impl_dx11.h"
#include "Injection/LoadLibrary.h"

#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "d3dcompiler.lib")
#pragma comment(lib, "dxgi.lib")
#pragma comment(lib, "comdlg32.lib")


#define MAX_LOADSTRING 100

// Data
static ID3D11Device*            g_pd3dDevice = nullptr;
static ID3D11DeviceContext*     g_pd3dDeviceContext = nullptr;
static IDXGISwapChain*          g_pSwapChain = nullptr;
static ID3D11RenderTargetView*  g_mainRenderTargetView = nullptr;


// Global Variables:
HINSTANCE hInst;                                // current instance
WCHAR szTitle[MAX_LOADSTRING];                  // The title bar text
WCHAR szWindowClass[MAX_LOADSTRING];            // the main window class name
static bool g_titleBarDragZone = false;         // set by ImGui each frame
static const float g_titleBarHeight = 32.0f;

// Callback for EnumWindows to collect PIDs that own visible windows
static BOOL CALLBACK EnumWindowsProcCollect(HWND hwnd, LPARAM lParam)
{
    if (IsWindowVisible(hwnd) && GetWindowTextLengthW(hwnd) > 0)
    {
        DWORD pid = 0;
        GetWindowThreadProcessId(hwnd, &pid);
        reinterpret_cast<std::set<DWORD>*>(lParam)->insert(pid);
    }
    return TRUE;
}

// Forward declarations of functions included in this code module:
ATOM                MyRegisterClass(HINSTANCE hInstance);
BOOL                InitInstance(HINSTANCE, int);
LRESULT CALLBACK    WndProc(HWND, UINT, WPARAM, LPARAM);
INT_PTR CALLBACK    About(HWND, UINT, WPARAM, LPARAM);
bool CreateDeviceD3D(HWND hWnd);
void CleanupDeviceD3D();
void CreateRenderTarget();
void CleanupRenderTarget();

int APIENTRY wWinMain(_In_ HINSTANCE hInstance,
                     _In_opt_ HINSTANCE hPrevInstance,
                     _In_ LPWSTR    lpCmdLine,
                     _In_ int       nCmdShow)
{
    UNREFERENCED_PARAMETER(hPrevInstance);
    UNREFERENCED_PARAMETER(lpCmdLine);

    // Initialize global strings
    LoadStringW(hInstance, IDS_APP_TITLE, szTitle, MAX_LOADSTRING);
    LoadStringW(hInstance, IDC_SIMPLEINJECTOR, szWindowClass, MAX_LOADSTRING);
    MyRegisterClass(hInstance);

    // Perform application initialization:
    if (!InitInstance (hInstance, nCmdShow))
    {
        return FALSE;
    }

    // Initialize Direct3D
    HWND hWnd = FindWindowW(szWindowClass, szTitle);
    hWnd = FindWindow(szWindowClass, szTitle);

    if (!CreateDeviceD3D(hWnd))
    {
        CleanupDeviceD3D();
        return 1;
    }

    // Setup Dear ImGui context
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGuiIO& io = ImGui::GetIO(); (void)io;
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;     // Enable Keyboard Controls

    // Setup Dear ImGui style
    ImGui::StyleColorsDark();

    // Setup Platform/Renderer backends
    ImGui_ImplWin32_Init(hWnd);
    ImGui_ImplDX11_Init(g_pd3dDevice, g_pd3dDeviceContext);

    // State for injection
    static char dllPath[MAX_PATH] = "";
    static int selectedProcessIndex = -1;
    struct ProcessInfo { DWORD id; std::string name; bool hasVisibleWindow; };
    static std::vector<ProcessInfo> processes;
    static bool processesLoaded = false;
    static bool showAppsOnly = false;
    static Injection::Result lastInjectionResult = Injection::Result::Success;
    static bool injectionAttempted = false;

    // Helper to refresh processes
    auto RefreshProcesses = [&]() {
        processes.clear();
        // Collect PIDs that own at least one visible window with a title
        std::set<DWORD> visiblePIDs;
        EnumWindows(EnumWindowsProcCollect, reinterpret_cast<LPARAM>(&visiblePIDs));

        HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if (hSnapshot != INVALID_HANDLE_VALUE) {
            PROCESSENTRY32 pe;
            pe.dwSize = sizeof(pe);
            if (Process32First(hSnapshot, &pe)) {
                do {
                    std::wstring wname = pe.szExeFile;
                    char nameBuf[256];
                    WideCharToMultiByte(CP_UTF8, 0, wname.c_str(), -1, nameBuf, sizeof(nameBuf), NULL, NULL);
                    bool hasWindow = visiblePIDs.count(pe.th32ProcessID) > 0;
                    processes.push_back({ pe.th32ProcessID, std::string(nameBuf), hasWindow });
                } while (Process32Next(hSnapshot, &pe));
            }
            CloseHandle(hSnapshot);
        }
    };

    if (!processesLoaded) { RefreshProcesses(); processesLoaded = true; }

    HACCEL hAccelTable = LoadAccelerators(hInstance, MAKEINTRESOURCE(IDC_SIMPLEINJECTOR));

    MSG msg;
    ZeroMemory(&msg, sizeof(msg));

    // Main message loop:
    while (msg.message != WM_QUIT)
    {
        if (PeekMessage(&msg, nullptr, 0, 0, PM_REMOVE))
        {
            if (!TranslateAccelerator(msg.hwnd, hAccelTable, &msg))
            {
                TranslateMessage(&msg);
                DispatchMessage(&msg);
            }
        }
        else
        {
            // Start the Dear ImGui frame
            ImGui_ImplDX11_NewFrame();
            ImGui_ImplWin32_NewFrame();
            ImGui::NewFrame();

            // UI
            {
                const ImGuiViewport* viewport = ImGui::GetMainViewport();
                ImGui::SetNextWindowPos(viewport->WorkPos);
                ImGui::SetNextWindowSize(viewport->WorkSize);
                ImGui::PushStyleVar(ImGuiStyleVar_WindowPadding, ImVec2(0, 0));
                ImGui::Begin("Simple Injector", nullptr, ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoResize | ImGuiWindowFlags_NoBringToFrontOnFocus);
                ImGui::PopStyleVar();

                // Custom title bar
                {
                    const float titleBarH = 32.0f;
                    const float btnW = 46.0f;
                    ImVec2 winPos = ImGui::GetWindowPos();
                    ImVec2 winSize = ImGui::GetWindowSize();
                    ImDrawList* dl = ImGui::GetWindowDrawList();

                    dl->AddRectFilled(winPos, ImVec2(winPos.x + winSize.x, winPos.y + titleBarH), IM_COL32(25, 25, 25, 255));
                    dl->AddText(ImVec2(winPos.x + 10.0f, winPos.y + 8.0f), IM_COL32(200, 200, 200, 255), "Simple Injector");
                    dl->AddLine(ImVec2(winPos.x, winPos.y + titleBarH), ImVec2(winPos.x + winSize.x, winPos.y + titleBarH), IM_COL32(60, 60, 60, 255));

                    // Close button
                    ImGui::SetCursorPos(ImVec2(winSize.x - btnW, 0.0f));
                    ImGui::PushStyleColor(ImGuiCol_Button, IM_COL32(25, 25, 25, 255));
                    ImGui::PushStyleColor(ImGuiCol_ButtonHovered, IM_COL32(200, 50, 50, 255));
                    ImGui::PushStyleColor(ImGuiCol_ButtonActive, IM_COL32(150, 30, 30, 255));
                    ImGui::PushStyleVar(ImGuiStyleVar_FrameRounding, 0.0f);
                    if (ImGui::Button("  X  ##close", ImVec2(btnW, titleBarH)))
                        PostMessage(hWnd, WM_CLOSE, 0, 0);
                    bool closeHovered = ImGui::IsItemHovered();
                    ImGui::PopStyleVar();
                    ImGui::PopStyleColor(3);

                    // Maximize / Restore button
                    ImGui::SetCursorPos(ImVec2(winSize.x - 2.0f * btnW, 0.0f));
                    ImGui::PushStyleColor(ImGuiCol_Button, IM_COL32(25, 25, 25, 255));
                    ImGui::PushStyleColor(ImGuiCol_ButtonHovered, IM_COL32(55, 55, 55, 255));
                    ImGui::PushStyleColor(ImGuiCol_ButtonActive, IM_COL32(75, 75, 75, 255));
                    ImGui::PushStyleVar(ImGuiStyleVar_FrameRounding, 0.0f);
                    WINDOWPLACEMENT wp; wp.length = sizeof(wp);
                    GetWindowPlacement(hWnd, &wp);
                    bool isMaximized = (wp.showCmd == SW_MAXIMIZE);
                    if (ImGui::Button(isMaximized ? " [] ##max" : " [] ##max", ImVec2(btnW, titleBarH)))
                        ShowWindow(hWnd, isMaximized ? SW_RESTORE : SW_MAXIMIZE);
                    bool maxHovered = ImGui::IsItemHovered();
                    ImGui::PopStyleVar();
                    ImGui::PopStyleColor(3);

                    // Minimize button
                    ImGui::SetCursorPos(ImVec2(winSize.x - 3.0f * btnW, 0.0f));
                    ImGui::PushStyleColor(ImGuiCol_Button, IM_COL32(25, 25, 25, 255));
                    ImGui::PushStyleColor(ImGuiCol_ButtonHovered, IM_COL32(55, 55, 55, 255));
                    ImGui::PushStyleColor(ImGuiCol_ButtonActive, IM_COL32(75, 75, 75, 255));
                    ImGui::PushStyleVar(ImGuiStyleVar_FrameRounding, 0.0f);
                    if (ImGui::Button("  _  ##min", ImVec2(btnW, titleBarH)))
                        ShowWindow(hWnd, SW_MINIMIZE);
                    bool minHovered = ImGui::IsItemHovered();
                    ImGui::PopStyleVar();
                    ImGui::PopStyleColor(3);

                    // Track drag zone for WM_NCHITTEST
                    ImVec2 mousePos = ImGui::GetIO().MousePos;
                    bool inTitleBar = mousePos.y >= winPos.y && mousePos.y <= winPos.y + titleBarH
                                  && mousePos.x >= winPos.x && mousePos.x <= winPos.x + winSize.x;
                    g_titleBarDragZone = inTitleBar && !closeHovered && !maxHovered && !minHovered;

                    ImGui::SetCursorPos(ImVec2(8.0f, titleBarH + 8.0f));
                }

                // Content area
                ImGui::BeginChild("##Content", ImVec2(ImGui::GetContentRegionAvail().x - 8.0f, ImGui::GetContentRegionAvail().y - 8.0f));

                if (ImGui::Button("Refresh Processes")) {
                    RefreshProcesses();
                    selectedProcessIndex = -1;
                }
                ImGui::SameLine();
                ImGui::Checkbox("Applications only", &showAppsOnly);

                ImGui::Text("Processes:");
                if (ImGui::BeginListBox("##ProcessList", ImVec2(-FLT_MIN, ImGui::GetContentRegionAvail().y - 80.0f)))
                {
                    for (int n = 0; n < (int)processes.size(); n++)
                    {
                        if (showAppsOnly && !processes[n].hasVisibleWindow)
                            continue;
                        ImGui::PushID(n);
                        bool is_selected = (selectedProcessIndex == n);
                        char label[512];
                        snprintf(label, sizeof(label), "%-40s  [%lu]", processes[n].name.c_str(), processes[n].id);
                        if (ImGui::Selectable(label, is_selected))
                            selectedProcessIndex = n;
                        if (is_selected)
                            ImGui::SetItemDefaultFocus();
                        ImGui::PopID();
                    }
                    ImGui::EndListBox();
                }

                ImGui::Text("DLL:");
                ImGui::SameLine();
                ImGui::SetNextItemWidth(ImGui::GetContentRegionAvail().x - 80.0f);
                ImGui::InputText("##DLLPath", dllPath, MAX_PATH, ImGuiInputTextFlags_ReadOnly);
                ImGui::SameLine();
                if (ImGui::Button("Browse", ImVec2(72.0f, 0.0f)))
                {
                    char tempPath[MAX_PATH];
                    strncpy_s(tempPath, dllPath, MAX_PATH);
                    OPENFILENAMEA ofn;
                    ZeroMemory(&ofn, sizeof(ofn));
                    ofn.lStructSize = sizeof(ofn);
                    ofn.hwndOwner = hWnd;
                    ofn.lpstrFilter = "DLL Files\0*.dll\0All Files\0*.*\0";
                    ofn.lpstrFile = tempPath;
                    ofn.nMaxFile = MAX_PATH;
                    ofn.Flags = OFN_FILEMUSTEXIST | OFN_PATHMUSTEXIST | OFN_NOCHANGEDIR;
                    if (GetOpenFileNameA(&ofn))
                        strncpy_s(dllPath, tempPath, MAX_PATH);
                }

                if (ImGui::Button("Inject"))
                {
                    if (selectedProcessIndex >= 0 && strlen(dllPath) > 0)
                    {
                        DWORD pid = processes[selectedProcessIndex].id;
                        lastInjectionResult = Injection::InjectLoadLibrary(pid, dllPath);
                        injectionAttempted = true;
                        ImGui::OpenPopup("InjectionResult");
                    }
                }

                if (ImGui::BeginPopupModal("InjectionResult", NULL, ImGuiWindowFlags_AlwaysAutoResize))
                {
                    bool success = (lastInjectionResult == Injection::Result::Success);
                    if (success)
                        ImGui::TextColored(ImVec4(0.4f, 1.0f, 0.4f, 1.0f), "%s", Injection::ResultToString(lastInjectionResult));
                    else
                        ImGui::TextColored(ImVec4(1.0f, 0.4f, 0.4f, 1.0f), "%s", Injection::ResultToString(lastInjectionResult));
                    if (ImGui::Button("OK", ImVec2(120, 0))) { ImGui::CloseCurrentPopup(); }
                    ImGui::EndPopup();
                }

                ImGui::EndChild();
                ImGui::End();
            }

            // Rendering
            ImGui::Render();
            g_pd3dDeviceContext->OMSetRenderTargets(1, &g_mainRenderTargetView, nullptr);
            static const float clear_color_with_alpha[4] = { 0.45f, 0.55f, 0.60f, 1.00f };
            g_pd3dDeviceContext->ClearRenderTargetView(g_mainRenderTargetView, clear_color_with_alpha);
            ImGui_ImplDX11_RenderDrawData(ImGui::GetDrawData());

            g_pSwapChain->Present(1, 0); // Present with vsync
        }
    }

    // Cleanup
    ImGui_ImplDX11_Shutdown();
    ImGui_ImplWin32_Shutdown();
    ImGui::DestroyContext();

    CleanupDeviceD3D();
    DestroyWindow(hWnd);
    UnregisterClassW(szWindowClass, hInstance);

    return (int) msg.wParam;
}



//
//  FUNCTION: MyRegisterClass()
//
//  PURPOSE: Registers the window class.
//
ATOM MyRegisterClass(HINSTANCE hInstance)
{
    WNDCLASSEXW wcex;

    wcex.cbSize = sizeof(WNDCLASSEX);

    wcex.style          = CS_HREDRAW | CS_VREDRAW;
    wcex.lpfnWndProc    = WndProc;
    wcex.cbClsExtra     = 0;
    wcex.cbWndExtra     = 0;
    wcex.hInstance      = hInstance;
    wcex.hIcon          = LoadIcon(hInstance, MAKEINTRESOURCE(IDI_SIMPLEINJECTOR));
    wcex.hCursor        = LoadCursor(nullptr, IDC_ARROW);
    wcex.hbrBackground  = (HBRUSH)(COLOR_WINDOW+1);
    wcex.lpszMenuName   = nullptr;
    wcex.lpszClassName  = szWindowClass;
    wcex.hIconSm        = LoadIcon(wcex.hInstance, MAKEINTRESOURCE(IDI_SMALL));

    return RegisterClassExW(&wcex);
}

//
//   FUNCTION: InitInstance(HINSTANCE, int)
//
//   PURPOSE: Saves instance handle and creates main window
//
//   COMMENTS:
//
//        In this function, we save the instance handle in a global variable and
//        create and display the main program window.
//
BOOL InitInstance(HINSTANCE hInstance, int nCmdShow)
{
   hInst = hInstance; // Store instance handle in our global variable

   int winW = 800, winH = 500;
   int posX = (GetSystemMetrics(SM_CXSCREEN) - winW) / 2;
   int posY = (GetSystemMetrics(SM_CYSCREEN) - winH) / 2;

   HWND hWnd = CreateWindowW(szWindowClass, szTitle, WS_OVERLAPPEDWINDOW,
      posX, posY, winW, winH, nullptr, nullptr, hInstance, nullptr);

   if (!hWnd)
   {
      return FALSE;
   }

   ShowWindow(hWnd, nCmdShow);
   UpdateWindow(hWnd);

   return TRUE;
}

//
//  FUNCTION: WndProc(HWND, UINT, WPARAM, LPARAM)
//
//  PURPOSE: Processes messages for the main window.
//
//  WM_COMMAND  - process the application menu
//  WM_PAINT    - Paint the main window
//  WM_DESTROY  - post a quit message and return
//
//
extern LRESULT ImGui_ImplWin32_WndProcHandler(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
{
    if (ImGui_ImplWin32_WndProcHandler(hWnd, message, wParam, lParam))
        return true;

    switch (message)
    {
    case WM_NCCALCSIZE:
        if (wParam == TRUE)
        {
            NCCALCSIZE_PARAMS* params = reinterpret_cast<NCCALCSIZE_PARAMS*>(lParam);
            HMONITOR monitor = MonitorFromWindow(hWnd, MONITOR_DEFAULTTONEAREST);
            MONITORINFO mi;
            mi.cbSize = sizeof(mi);
            if (GetMonitorInfo(monitor, &mi))
            {
                WINDOWPLACEMENT wp;
                wp.length = sizeof(wp);
                GetWindowPlacement(hWnd, &wp);
                if (wp.showCmd == SW_MAXIMIZE)
                    params->rgrc[0] = mi.rcWork;
            }
        }
        return 0;
    case WM_NCHITTEST:
    {
        POINT pt;
        pt.x = (int)(short)LOWORD(lParam);
        pt.y = (int)(short)HIWORD(lParam);
        ScreenToClient(hWnd, &pt);
        RECT rc;
        GetClientRect(hWnd, &rc);
        const int border = 6;
        if (pt.y <= border) {
            if (pt.x <= border) return HTTOPLEFT;
            if (pt.x >= rc.right - border) return HTTOPRIGHT;
            return HTTOP;
        }
        if (pt.y >= rc.bottom - border) {
            if (pt.x <= border) return HTBOTTOMLEFT;
            if (pt.x >= rc.right - border) return HTBOTTOMRIGHT;
            return HTBOTTOM;
        }
        if (pt.x <= border) return HTLEFT;
        if (pt.x >= rc.right - border) return HTRIGHT;
        if (pt.y <= (int)g_titleBarHeight && g_titleBarDragZone)
            return HTCAPTION;
        return HTCLIENT;
    }
    case WM_GETMINMAXINFO:
    {
        MINMAXINFO* mmi = reinterpret_cast<MINMAXINFO*>(lParam);
        mmi->ptMinTrackSize.x = 500;
        mmi->ptMinTrackSize.y = 350;
        return 0;
    }
    case WM_SIZE:
        if (g_pd3dDevice != nullptr && wParam != SIZE_MINIMIZED)
        {
            CleanupRenderTarget();
            g_pSwapChain->ResizeBuffers(0, (UINT)LOWORD(lParam), (UINT)HIWORD(lParam), DXGI_FORMAT_UNKNOWN, 0);
            CreateRenderTarget();
        }
        return 0;
    case WM_SYSCOMMAND:
        if ((wParam & 0xfff0) == SC_KEYMENU)
            return 0;
        break;
    case WM_DESTROY:
        PostQuitMessage(0);
        return 0;
    }
    return DefWindowProc(hWnd, message, wParam, lParam);
}

// Helper functions

bool CreateDeviceD3D(HWND hWnd)
{
    DXGI_SWAP_CHAIN_DESC sd;
    ZeroMemory(&sd, sizeof(sd));
    sd.BufferCount = 2;
    sd.BufferDesc.Width = 0;
    sd.BufferDesc.Height = 0;
    sd.BufferDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    sd.BufferDesc.RefreshRate.Numerator = 60;
    sd.BufferDesc.RefreshRate.Denominator = 1;
    sd.Flags = DXGI_SWAP_CHAIN_FLAG_ALLOW_MODE_SWITCH;
    sd.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
    sd.OutputWindow = hWnd;
    sd.SampleDesc.Count = 1;
    sd.SampleDesc.Quality = 0;
    sd.Windowed = TRUE;
    sd.SwapEffect = DXGI_SWAP_EFFECT_DISCARD;

    UINT createDeviceFlags = 0;
    D3D_FEATURE_LEVEL featureLevel;
    const D3D_FEATURE_LEVEL featureLevelArray[2] = { D3D_FEATURE_LEVEL_11_0, D3D_FEATURE_LEVEL_10_0, };
    if (D3D11CreateDeviceAndSwapChain(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, createDeviceFlags, featureLevelArray, 2, D3D11_SDK_VERSION, &sd, &g_pSwapChain, &g_pd3dDevice, &featureLevel, &g_pd3dDeviceContext) != S_OK)
        return false;

    CreateRenderTarget();
    return true;
}

void CleanupDeviceD3D()
{
    CleanupRenderTarget();
    if (g_pSwapChain) { g_pSwapChain->Release(); g_pSwapChain = nullptr; }
    if (g_pd3dDeviceContext) { g_pd3dDeviceContext->Release(); g_pd3dDeviceContext = nullptr; }
    if (g_pd3dDevice) { g_pd3dDevice->Release(); g_pd3dDevice = nullptr; }
}

void CreateRenderTarget()
{
    ID3D11Texture2D* pBackBuffer;
    g_pSwapChain->GetBuffer(0, IID_PPV_ARGS(&pBackBuffer));
    g_pd3dDevice->CreateRenderTargetView(pBackBuffer, nullptr, &g_mainRenderTargetView);
    pBackBuffer->Release();
}

void CleanupRenderTarget()
{
    if (g_mainRenderTargetView) { g_mainRenderTargetView->Release(); g_mainRenderTargetView = nullptr; }
}