Inject multiple DLLs into a process

[QQ2017]

MapFromMemory cannot inject multiple DLLs into a process. There is no problem injecting a single DLL into the process.

#include <iostream>
#include <set>
#include <string>
#include "BlackBone/Process/Process.h"
#include "BlackBone/Misc/DynImport.h"

using namespace blackbone;

// 定义 DLL 加载数据结构体
struct DllLoadData
{
    bool initialized = false;
    bool deinitialized = false;
    bool static_initialized = false;
    bool static_thread_initialized = false;
    bool static_func_initialized = false;
    bool seh_internal = false;
    bool seh_external = false;
    bool ceh_internal = false;
    bool ceh_external = false;
    bool ceh_uncaught = false;
};

// 定义用于回调的模块集合
std::set<std::wstring> nativeMods = { L"combase.dll", L"user32.dll" };
std::set<std::wstring> modList = { L"windows.storage.dll", L"shell32.dll", L"shlwapi.dll" };


// 从文件读取数据
std::pair<std::unique_ptr<uint8_t[]>, uint32_t> GetFileData(const std::wstring& path)
{
    HANDLE hFile = CreateFileW(path.c_str(), FILE_GENERIC_READ, FILE_SHARE_READ, nullptr, OPEN_EXISTING, 0, nullptr);
    if (hFile == INVALID_HANDLE_VALUE)
    {
        std::wcerr << L"Failed to open file: " << path << std::endl;
        return std::make_pair(nullptr, 0);
    }

    uint32_t size = GetFileSize(hFile, nullptr);
    auto buf = std::make_unique<uint8_t[]>(size);

    DWORD bytes = 0;
    if (!ReadFile(hFile, buf.get(), size, &bytes, nullptr) || bytes != size)
    {
        std::wcerr << L"Failed to read file: " << path << std::endl;
        CloseHandle(hFile);
        return std::make_pair(nullptr, 0);
    }

    CloseHandle(hFile);
    return std::make_pair(std::move(buf), size);
}

// DLL 从文件手动映射函数
void MapFromFile(const std::wstring& hostPath, const std::wstring& dllPath)
{
    Process proc;
    NTSTATUS status = proc.Attach(hostPath.c_str());
    if (!NT_SUCCESS(status))
    {
        std::wcerr << L"Failed to attach to process: " << hostPath << std::endl;
        return;
    }
    proc.EnsureInit();

    // 内置回调函数
    auto MapCallback = [](CallbackType type, void* /*context*/, Process& /*process*/, const ModuleData& modInfo) -> LoadData
        {
            static std::set<std::wstring> nativeMods = { L"combase.dll", L"user32.dll" };
            static std::set<std::wstring> modList = { L"windows.storage.dll", L"shell32.dll", L"shlwapi.dll" };

            if (type == PreCallback)
            {
                // 跳过系统模块的加载行为，只做映射
                if (nativeMods.count(modInfo.name))
                {
                    std::wcout << L"Skipping loading native module: " << modInfo.name << std::endl;
                    return LoadData(MT_Native, Ldr_None);
                }
            }
            else if (type == PostCallback)
            {
                // 对于 modList 中的模块，执行标准加载行为
                if (modList.count(modInfo.name))
                {
                    std::wcout << L"Loading module into Ldr: " << modInfo.name << std::endl;
                    return LoadData(MT_Default, Ldr_ModList);
                }
            }

            return LoadData(MT_Default, Ldr_None);  // 默认行为
        };

    auto image = proc.mmap().MapImage(dllPath, NoFlags, MapCallback);
    if (!image.success())
    {
        std::wcerr << L"Failed to map image: " << dllPath << std::endl;
        return;
    }

    std::wcout << L"Successfully mapped image from file: " << dllPath << std::endl;

    //auto g_loadDataPtr = proc.modules().GetExport(image.result(), "g_LoadData");
    //if (g_loadDataPtr.success())
    //{
    //    auto g_loadData = proc.memory().Read<DllLoadData>(g_loadDataPtr->procAddress);
    //    if (g_loadData.success())
    //    {
    //        std::wcout << L"g_LoadData successfully read from process memory" << std::endl;
    //    }
    //}

    proc.Detach();
}

// DLL 从内存手动映射函数
void MapFromMemory(const std::wstring& hostPath, const std::wstring& dllPath)
{
    // 读取 DLL 文件到内存
    auto [buf, size] = GetFileData(dllPath);
    if (size == 0)
    {
        std::wcerr << L"Failed to read DLL file: " << dllPath << std::endl;
        return;
    }

    Process proc;
    NTSTATUS status = proc.Attach(hostPath.c_str());
    if (!NT_SUCCESS(status))
    {
        std::wcerr << L"Failed to attach to process: " << hostPath << std::endl;
        return;
    }
    proc.EnsureInit();


    // 内置回调函数
    auto MapCallback = [](CallbackType type, void* /*context*/, Process& /*process*/, const ModuleData& modInfo) -> LoadData
        {
            static std::set<std::wstring> nativeMods = { L"combase.dll", L"user32.dll" };
            static std::set<std::wstring> modList = { L"windows.storage.dll", L"shell32.dll", L"shlwapi.dll" };

            if (type == PreCallback)
            {
                // 跳过系统模块的加载行为，只做映射
                if (nativeMods.count(modInfo.name))
                {
                    std::wcout << L"Skipping loading native module: " << modInfo.name << std::endl;
                    return LoadData(MT_Native, Ldr_None);
                }
            }
            else if (type == PostCallback)
            {
                // 对于 modList 中的模块，执行标准加载行为
                if (modList.count(modInfo.name))
                {
                    std::wcout << L"Loading module into Ldr: " << modInfo.name << std::endl;
                    return LoadData(MT_Default, Ldr_ModList);
                }
            }

            return LoadData(MT_Default, Ldr_None);  // 默认行为
        };


    // 将 DLL 文件从内存映射到目标进程
    auto image = proc.mmap().MapImage(size, buf.get(), false);
    if (!image.success())
    {
        std::wcerr << L"Failed to map image from memory" << std::endl;
        std::wcout << L"Mapping failed with error 0x" << std::hex << image.status
            << L". " << Utils::GetErrorDescription(image.status) << std::endl << std::endl;
        
        return;
    }

    std::wcout << L"Successfully mapped image from memory: " << dllPath << std::endl;

    //auto g_loadDataPtr = proc.modules().GetExport(image.result(), "g_LoadData");
    //if (g_loadDataPtr.success())
    //{
    //    auto g_loadData = proc.memory().Read<DllLoadData>(g_loadDataPtr->procAddress);
    //    if (g_loadData.success())
    //    {
    //        std::wcout << L"g_LoadData successfully read from process memory" << std::endl;
    //    }
    //}

    proc.Detach();
}

// 主程序入口
int wmain(int argc, wchar_t* argv[])
{
    if (argc < 4)
    {
        std::wcerr << L"Usage: <inject_method: file/memory> <target_process_path> <dll_path>" << std::endl;
        return -1;
    }
 
    std::wstring method = argv[1];
    std::wstring targetProcessPath = argv[2];
    std::wstring dllPath = argv[3];

    if (method == L"file")
    {
        std::wcout << L"Injecting DLL from file: " << dllPath << L" into process: " << targetProcessPath << std::endl;
        MapFromFile(targetProcessPath, dllPath);
    }
    else if (method == L"memory")
    {
        std::wcout << L"Injecting DLL from memory: " << dllPath << L" into process: " << targetProcessPath << std::endl;
        MapFromMemory(targetProcessPath, dllPath);
    }
    else
    {
        std::wcerr << L"Unknown method: " << method << std::endl;
        return -1;
    }

    return 0;
}