#include "sandbox.h"


auto Sandbox::InitializeLdrData() -> void {
    if (m_peInfo->isX64) {
        // 为LDR_DATA分配内存
        uint64_t ldrDataAddress = m_pebBase + sizeof(X64PEB);
        m_pebEnd = ldrDataAddress + sizeof(X64_PEB_LDR_DATA);
        m_peb64.Ldr = ldrDataAddress;

        // 映射LDR数据内存
        uc_mem_map(m_ucEngine, ldrDataAddress, sizeof(X64_PEB_LDR_DATA),
            UC_PROT_ALL);

        // 初始化LDR_DATA结构
        X64_PEB_LDR_DATA ldrData = { 0 };
        ldrData.Length = sizeof(X64_PEB_LDR_DATA);
        ldrData.Initialized = 1;

        // 初始化链表头 - 每个链表头都指向自己(空链表)
        uint64_t inLoadOrderListHead = ldrDataAddress + offsetof(X64_PEB_LDR_DATA, InLoadOrderModuleList);
        uint64_t inMemoryOrderListHead = ldrDataAddress + offsetof(X64_PEB_LDR_DATA, InMemoryOrderModuleList);
        uint64_t inInitOrderListHead = ldrDataAddress + offsetof(X64_PEB_LDR_DATA, InInitializationOrderModuleList);

        // 设置每个链表头的Flink和Blink都指向自己
        ldrData.InLoadOrderModuleList.Flink = (LIST_ENTRY*)inLoadOrderListHead;
        ldrData.InLoadOrderModuleList.Blink = (LIST_ENTRY*)inLoadOrderListHead;

        ldrData.InMemoryOrderModuleList.Flink = (LIST_ENTRY*)inMemoryOrderListHead;
        ldrData.InMemoryOrderModuleList.Blink = (LIST_ENTRY*)inMemoryOrderListHead;

        ldrData.InInitializationOrderModuleList.Flink = (LIST_ENTRY*)inInitOrderListHead;
        ldrData.InInitializationOrderModuleList.Blink = (LIST_ENTRY*)inInitOrderListHead;

        // 写入LDR_DATA结构到内存
        uc_mem_write(m_ucEngine, ldrDataAddress, &ldrData, sizeof(X64_PEB_LDR_DATA));

        // 更新PEB中的Ldr指针
        uc_mem_write(m_ucEngine, m_pebBase, &m_peb64, sizeof(X64PEB));
    }
}
auto Sandbox::CreateLdrEntry(const std::shared_ptr<struct_moudle>& module,
    uint64_t entryAddress, uint64_t fullNameAddress,
    uint64_t baseNameAddress) -> LDR_DATA_TABLE_ENTRY {

    // 清空整个结构体以确保所有字段正确初始化
    LDR_DATA_TABLE_ENTRY entry = {};
    if (LOG_LEVEL > 4) {

        // 检查模块基址 - 避免null指针
        if (module->base == 0) {
            printf("警告：模块'%s'的基址为0，这可能导致问题\n", module->name);
        }
    }

    entry.DllBase = reinterpret_cast<PVOID>(module->base);
    entry.EntryPoint = reinterpret_cast<PVOID>(module->base + module->entry);
    entry.SizeOfImages = static_cast<ULONG>(module->size);
    // 准备模块名称的Unicode字符串
    wchar_t nameBuffer[MAX_PATH] = { 0 };
    size_t convertedChars = 0;
    std::string fakeFullName = "c:\\windows\\system32\\" + std::string(module->name);

    // 使用更安全的mbstowcs_s函数
    mbstowcs_s(&convertedChars, nameBuffer, MAX_PATH, fakeFullName.c_str(), fakeFullName.size());

    // 验证转换后的字符串长度
    size_t nameLen = wcslen(nameBuffer);
    if (LOG_LEVEL > 4) {

        if (nameLen == 0 || nameLen >= MAX_PATH) {
            printf("警告：模块名称'%s'转换问题，长度=%zu\n", module->name, nameLen);
        }
    }
    // 设置字符串
    entry.FullDllName.Length = static_cast<USHORT>(nameLen * sizeof(wchar_t));
    entry.FullDllName.MaximumLength = MAX_PATH * sizeof(wchar_t);
    entry.FullDllName.Buffer = reinterpret_cast<PWSTR>(fullNameAddress);

    entry.BaseDllName.Length = static_cast<USHORT>(nameLen * sizeof(wchar_t));
    entry.BaseDllName.MaximumLength = MAX_PATH * sizeof(wchar_t);
    entry.BaseDllName.Buffer = reinterpret_cast<PWSTR>(baseNameAddress);
    if (LOG_LEVEL > 4) {
        // 写入字符串数据 - 添加调试输出
        printf("DEBUG: 写入模块名 '%s' 到地址 0x%llx 和 0x%llx\n",
            module->name, fullNameAddress, baseNameAddress);
    }
    uc_mem_write(m_ucEngine, fullNameAddress, nameBuffer, (nameLen + 1) * sizeof(wchar_t));
    uc_mem_write(m_ucEngine, baseNameAddress, nameBuffer, (nameLen + 1) * sizeof(wchar_t));

    // 初始化链表指针为NULL（将在UpdateLdrLinks中设置）
    entry.InLoadOrderLinks.Flink = nullptr;
    entry.InLoadOrderLinks.Blink = nullptr;
    entry.InMemoryOrderLinks.Flink = nullptr;
    entry.InMemoryOrderLinks.Blink = nullptr;
    entry.InInitializationOrderLinks.Flink = nullptr;
    entry.InInitializationOrderLinks.Blink = nullptr;

    return entry;
}
auto Sandbox::UpdateLdrLinks(const LDR_DATA_TABLE_ENTRY& entry,
    uint64_t entryAddress, X64_PEB_LDR_DATA& ldrData) -> void {
    // 获取链表头地址
    uint64_t inLoadOrderListHead = m_peb64.Ldr + offsetof(X64_PEB_LDR_DATA, InLoadOrderModuleList);
    uint64_t inMemoryOrderListHead = m_peb64.Ldr + offsetof(X64_PEB_LDR_DATA, InMemoryOrderModuleList);
    uint64_t inInitOrderListHead = m_peb64.Ldr + offsetof(X64_PEB_LDR_DATA, InInitializationOrderModuleList);

    // 新条目在每个链表中的地址
    uint64_t entryInLoadOrder = entryAddress + offsetof(LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
    uint64_t entryInMemoryOrder = entryAddress + offsetof(LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks);
    uint64_t entryInInitOrder = entryAddress + offsetof(LDR_DATA_TABLE_ENTRY, InInitializationOrderLinks);

    LIST_ENTRY listHead;
    LIST_ENTRY entryLinks;

    // 处理InLoadOrderModuleList
    uc_mem_read(m_ucEngine, inLoadOrderListHead, &listHead, sizeof(LIST_ENTRY));

    // 更新条目的链接 (新条目插入到链表头之后)
    entryLinks.Flink = listHead.Flink;  // 新条目的Flink = 原来头的Flink
    entryLinks.Blink = (LIST_ENTRY*)inLoadOrderListHead;  // 新条目的Blink = 头
    uc_mem_write(m_ucEngine, entryInLoadOrder, &entryLinks, sizeof(LIST_ENTRY));

    // 更新链表头的Flink指向新条目
    listHead.Flink = (LIST_ENTRY*)entryInLoadOrder;
    uc_mem_write(m_ucEngine, inLoadOrderListHead, &listHead, sizeof(LIST_ENTRY));

    // 如果不是空链表，更新原来第一个节点的Blink
    if (entryLinks.Flink != (LIST_ENTRY*)inLoadOrderListHead) {
        LIST_ENTRY nextEntry;
        uc_mem_read(m_ucEngine, (uint64_t)entryLinks.Flink, &nextEntry, sizeof(LIST_ENTRY));
        nextEntry.Blink = (LIST_ENTRY*)entryInLoadOrder;
        uc_mem_write(m_ucEngine, (uint64_t)entryLinks.Flink, &nextEntry, sizeof(LIST_ENTRY));
    }

    // 同理处理InMemoryOrderModuleList
    uc_mem_read(m_ucEngine, inMemoryOrderListHead, &listHead, sizeof(LIST_ENTRY));
    entryLinks.Flink = listHead.Flink;
    entryLinks.Blink = (LIST_ENTRY*)inMemoryOrderListHead;
    uc_mem_write(m_ucEngine, entryInMemoryOrder, &entryLinks, sizeof(LIST_ENTRY));
    listHead.Flink = (LIST_ENTRY*)entryInMemoryOrder;
    uc_mem_write(m_ucEngine, inMemoryOrderListHead, &listHead, sizeof(LIST_ENTRY));
    if (entryLinks.Flink != (LIST_ENTRY*)inMemoryOrderListHead) {
        LIST_ENTRY nextEntry;
        uc_mem_read(m_ucEngine, (uint64_t)entryLinks.Flink, &nextEntry, sizeof(LIST_ENTRY));
        nextEntry.Blink = (LIST_ENTRY*)entryInMemoryOrder;
        uc_mem_write(m_ucEngine, (uint64_t)entryLinks.Flink, &nextEntry, sizeof(LIST_ENTRY));
    }

    // 处理InInitializationOrderModuleList
    uc_mem_read(m_ucEngine, inInitOrderListHead, &listHead, sizeof(LIST_ENTRY));
    entryLinks.Flink = listHead.Flink;
    entryLinks.Blink = (LIST_ENTRY*)inInitOrderListHead;
    uc_mem_write(m_ucEngine, entryInInitOrder, &entryLinks, sizeof(LIST_ENTRY));
    listHead.Flink = (LIST_ENTRY*)entryInInitOrder;
    uc_mem_write(m_ucEngine, inInitOrderListHead, &listHead, sizeof(LIST_ENTRY));
    if (entryLinks.Flink != (LIST_ENTRY*)inInitOrderListHead) {
        LIST_ENTRY nextEntry;
        uc_mem_read(m_ucEngine, (uint64_t)entryLinks.Flink, &nextEntry, sizeof(LIST_ENTRY));
        nextEntry.Blink = (LIST_ENTRY*)entryInInitOrder;
        uc_mem_write(m_ucEngine, (uint64_t)entryLinks.Flink, &nextEntry, sizeof(LIST_ENTRY));
    }
}
auto Sandbox::AddModuleToLdr(const std::shared_ptr<struct_moudle>& module) -> void {
    if (!m_peInfo->isX64) {
        return;
    }
    if (LOG_LEVEL > 4) {
        // 调试输出
        printf("添加模块: '%s', 基址: 0x%llx, 入口点: 0x%llx, 大小: 0x%llx\n",
            module->name, module->base, module->entry, module->size);
    }
    // 计算所需内存
    const size_t stringSectionSize = MAX_PATH * sizeof(wchar_t) * 2;
    uint64_t entrySize = sizeof(LDR_DATA_TABLE_ENTRY) + stringSectionSize;

    // 确保地址对齐到8字节边界（对于64位结构体很重要）
    uint64_t entryAddress = (m_pebEnd + 7) & ~7ULL;
    if (LOG_LEVEL > 4) {
        // 调试输出
        printf("分配LDR条目: 地址=0x%llx, 大小=0x%llx\n", entryAddress, entrySize);
    }
    // 映射内存
    const size_t pageSize = 4096;
    size_t alignedSize = (entrySize + pageSize - 1) & ~(pageSize - 1);

    uc_err err = uc_mem_map(m_ucEngine, entryAddress, alignedSize, UC_PROT_ALL);
    if (err != UC_ERR_OK) {
        if (LOG_LEVEL > 4) {
            printf("错误: 无法映射LDR条目内存 @ 0x%llx (错误=%d)\n", entryAddress, err);
        }
        return;
    }

    // 设置Unicode字符串地址，确保8字节对齐
    uint64_t fullNameAddress = (entryAddress + sizeof(LDR_DATA_TABLE_ENTRY) + 7) & ~7ULL;
    uint64_t baseNameAddress = (fullNameAddress + (MAX_PATH * sizeof(wchar_t)) + 7) & ~7ULL;
    if (LOG_LEVEL > 4) {
        // 验证地址范围
        printf("调试: 字符串地址 - FullName=0x%llx, BaseName=0x%llx\n",
            fullNameAddress, baseNameAddress);
    }
    // 确保字符串地址不超出分配的内存范围
    if (baseNameAddress + (MAX_PATH * sizeof(wchar_t)) > entryAddress + alignedSize) {
        if (LOG_LEVEL > 4) {
            printf("错误: 字符串缓冲区超出分配的内存范围!\n");
        }
        return;
    }

    // 创建LDR条目
    auto entry = CreateLdrEntry(module, entryAddress, fullNameAddress, baseNameAddress);
    if (LOG_LEVEL > 4) {

        // 验证DllBase设置正确
        if ((uint64_t)entry.DllBase != module->base) {
            printf("警告: DllBase不匹配 (设置=%llx, 期望=%llx)\n",
                (uint64_t)entry.DllBase, module->base);
        }
    }
    // 从PEB读取当前LDR_DATA结构
    X64_PEB_LDR_DATA ldrData;
    uc_mem_read(m_ucEngine, m_peb64.Ldr, &ldrData, sizeof(X64_PEB_LDR_DATA));

    // 写入LDR_DATA_TABLE_ENTRY结构
    uc_mem_write(m_ucEngine, entryAddress, &entry, sizeof(LDR_DATA_TABLE_ENTRY));

    // 更新链表
    UpdateLdrLinks(entry, entryAddress, ldrData);

    // 更新PEB尾部位置
    m_pebEnd = entryAddress + alignedSize;
    if (LOG_LEVEL > 4) {

        printf("成功添加模块 '%s' 到LDR数据表 @ 0x%llx\n", module->name, entryAddress);
    }
}
auto Sandbox::TestLdrListTraversal() -> bool {
    if (LOG_LEVEL > 4) {

        printf("\n============ LDR链表遍历测试 ============\n");

        // 1. 读取PEB地址
        uint64_t pebAddress = m_pebBase; // 填入你的PEB基址
        X64PEB peb;
        if (uc_mem_read(m_ucEngine, pebAddress, &peb, sizeof(X64PEB)) != UC_ERR_OK) {
            printf("ERROR: 无法读取PEB结构体\n");
            return false;
        }

        printf("PEB位于: 0x%llx\n", pebAddress);
        printf("PEB.Ldr指向: 0x%llx\n", peb.Ldr);

        // 2. 读取LDR_DATA结构
        if (peb.Ldr == 0) {
            printf("ERROR: PEB.Ldr为NULL，LDR未初始化\n");
            return false;
        }

        X64_PEB_LDR_DATA ldrData;
        if (uc_mem_read(m_ucEngine, peb.Ldr, &ldrData, sizeof(X64_PEB_LDR_DATA)) != UC_ERR_OK) {
            printf("ERROR: 无法读取LDR_DATA结构体\n");
            return false;
        }

        printf("LDR_DATA.Length: %u\n", ldrData.Length);
        printf("LDR_DATA.Initialized: %u\n", ldrData.Initialized);

        // 3. 检查并打印所有链表
        DumpLdrList("InLoadOrderModuleList", peb.Ldr,
            offsetof(X64_PEB_LDR_DATA, InLoadOrderModuleList),
            offsetof(LDR_DATA_TABLE_ENTRY, InLoadOrderLinks));

        DumpLdrList("InMemoryOrderModuleList", peb.Ldr,
            offsetof(X64_PEB_LDR_DATA, InMemoryOrderModuleList),
            offsetof(LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks));

        DumpLdrList("InInitializationOrderModuleList", peb.Ldr,
            offsetof(X64_PEB_LDR_DATA, InInitializationOrderModuleList),
            offsetof(LDR_DATA_TABLE_ENTRY, InInitializationOrderLinks));

        printf("========================================\n");
    }
    return true;
}
auto Sandbox::FinalizeLdrLinks() -> void {
    if (!m_peInfo->isX64) {
        return;
    }

    // 获取链表头地址
    uint64_t inLoadOrderListHead = m_peb64.Ldr + offsetof(X64_PEB_LDR_DATA, InLoadOrderModuleList);
    uint64_t inMemoryOrderListHead = m_peb64.Ldr + offsetof(X64_PEB_LDR_DATA, InMemoryOrderModuleList);
    uint64_t inInitOrderListHead = m_peb64.Ldr + offsetof(X64_PEB_LDR_DATA, InInitializationOrderModuleList);

    // 闭合InLoadOrderModuleList链表
    CloseLdrList(inLoadOrderListHead, offsetof(LDR_DATA_TABLE_ENTRY, InLoadOrderLinks));

    // 闭合InMemoryOrderModuleList链表
    CloseLdrList(inMemoryOrderListHead, offsetof(LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks));

    // 闭合InInitializationOrderModuleList链表
    CloseLdrList(inInitOrderListHead, offsetof(LDR_DATA_TABLE_ENTRY, InInitializationOrderLinks));
}

auto Sandbox::CloseLdrList(uint64_t listHeadAddr, size_t entryLinkOffset) -> void {
    LIST_ENTRY listHead;
    uc_mem_read(m_ucEngine, listHeadAddr, &listHead, sizeof(LIST_ENTRY));

    // 如果链表为空，无需处理
    if (listHead.Flink == (LIST_ENTRY*)listHeadAddr) {
        return;
    }

    // 找到链表的最后一个元素
    uint64_t currentLink = (uint64_t)listHead.Flink;
    uint64_t lastLink = 0;

    while (currentLink != listHeadAddr && currentLink != 0) {
        LIST_ENTRY currentEntry;
        uc_mem_read(m_ucEngine, currentLink, &currentEntry, sizeof(LIST_ENTRY));

        // 如果下一个链接为空或指向自己，则当前节点是最后一个有效节点
        if (currentEntry.Flink == nullptr ||
            (uint64_t)currentEntry.Flink == currentLink ||
            (uint64_t)currentEntry.Flink == 0) {
            lastLink = currentLink;
            break;
        }

        // 如果下一个链接是链表头，说明链表已正确闭合
        if ((uint64_t)currentEntry.Flink == listHeadAddr) {
            return; // 链表已正确闭合，无需修复
        }

        lastLink = currentLink;
        currentLink = (uint64_t)currentEntry.Flink;
    }

    // 如果找到最后一个元素，则修复链接
    if (lastLink != 0) {
        // 读取最后一个元素
        LIST_ENTRY lastEntry;
        uc_mem_read(m_ucEngine, lastLink, &lastEntry, sizeof(LIST_ENTRY));

        // 设置最后一个元素的Flink指向链表头
        lastEntry.Flink = (LIST_ENTRY*)listHeadAddr;
        uc_mem_write(m_ucEngine, lastLink, &lastEntry, sizeof(LIST_ENTRY));

        // 设置链表头的Blink指向最后一个元素
        listHead.Blink = (LIST_ENTRY*)lastLink;
        uc_mem_write(m_ucEngine, listHeadAddr, &listHead, sizeof(LIST_ENTRY));

        if (LOG_LEVEL > 4) {
            printf("已修复链表：链表头=0x%llx, 最后元素=0x%llx\n",
                listHeadAddr, lastLink);
        }
    }
}
// 辅助函数：遍历并打印指定的LDR链表
auto Sandbox::DumpLdrList(const char* listName, uint64_t ldrDataBase, size_t listOffset, size_t entryLinkOffset) -> void {
    if (LOG_LEVEL > 4) {

        printf("\n--- %s 链表内容 ---\n", listName);

        // 获取链表头地址
        uint64_t listHeadAddr = ldrDataBase + listOffset;
        LIST_ENTRY listHead;
        if (uc_mem_read(m_ucEngine, listHeadAddr, &listHead, sizeof(LIST_ENTRY)) != UC_ERR_OK) {
            printf("ERROR: 无法读取链表头 @ 0x%llx\n", listHeadAddr);
            return;
        }

        printf("列表头 @ 0x%llx: Flink=0x%llx, Blink=0x%llx\n",
            listHeadAddr, (uint64_t)listHead.Flink, (uint64_t)listHead.Blink);

        // 检查链表是否为空
        if (listHead.Flink == (LIST_ENTRY*)listHeadAddr) {
            printf("链表为空 (Flink指向自身)\n");
            return;
        }

        // 遍历链表
        uint64_t currentLink = (uint64_t)listHead.Flink;
        int entryCount = 0;

        while (currentLink != listHeadAddr && entryCount < 100) { // 防止无限循环
            entryCount++;

            // 计算当前LDR_DATA_TABLE_ENTRY的地址
            uint64_t entryAddr = currentLink - entryLinkOffset;

            // 读取整个条目
            LDR_DATA_TABLE_ENTRY entry;
            if (uc_mem_read(m_ucEngine, entryAddr, &entry, sizeof(LDR_DATA_TABLE_ENTRY)) != UC_ERR_OK) {
                printf("ERROR: 无法读取条目 @ 0x%llx\n", entryAddr);
                break;
            }

            printf("\n条目 #%d @ 0x%llx:\n", entryCount, entryAddr);
            printf("  DllBase: 0x%llx\n", (uint64_t)entry.DllBase);
            printf("  EntryPoint: 0x%llx\n", (uint64_t)entry.EntryPoint);
            printf("  SizeOfImage: 0x%x\n", entry.SizeOfImages);

            // 读取和打印模块名称
            if ((uint64_t)entry.FullDllName.Buffer != 0 && entry.FullDllName.Length > 0) {
                size_t nameLen = entry.FullDllName.Length / sizeof(wchar_t);
                wchar_t* nameBuffer = new wchar_t[nameLen + 1];

                if (uc_mem_read(m_ucEngine, (uint64_t)entry.FullDllName.Buffer, nameBuffer, entry.FullDllName.Length) == UC_ERR_OK) {
                    nameBuffer[nameLen] = L'\0';

                    // 转换为多字节用于输出
                    char mbName[MAX_PATH * 2] = { 0 };
                    wcstombs(mbName, nameBuffer, sizeof(mbName) - 1);
                    printf("  FullDllName: %s\n", mbName);
                }
                else {
                    printf("  ERROR: 无法读取FullDllName @ 0x%llx (长度: %d)\n",
                        (uint64_t)entry.FullDllName.Buffer, entry.FullDllName.Length);
                }

                delete[] nameBuffer;
            }
            else {
                printf("  FullDllName: <空或无效>\n");
            }

            // 同样打印BaseDllName
            if ((uint64_t)entry.BaseDllName.Buffer != 0 && entry.BaseDllName.Length > 0) {
                size_t nameLen = entry.BaseDllName.Length / sizeof(wchar_t);
                wchar_t* nameBuffer = new wchar_t[nameLen + 1];

                if (uc_mem_read(m_ucEngine, (uint64_t)entry.BaseDllName.Buffer, nameBuffer, entry.BaseDllName.Length) == UC_ERR_OK) {
                    nameBuffer[nameLen] = L'\0';

                    char mbName[MAX_PATH * 2] = { 0 };
                    wcstombs(mbName, nameBuffer, sizeof(mbName) - 1);
                    printf("  BaseDllName: %s\n", mbName);
                }
                else {
                    printf("  ERROR: 无法读取BaseDllName @ 0x%llx (长度: %d)\n",
                        (uint64_t)entry.BaseDllName.Buffer, entry.BaseDllName.Length);
                }

                delete[] nameBuffer;
            }
            else {
                printf("  BaseDllName: <空或无效>\n");
            }

            // 检查链接指针的有效性
            printf("  链接: Flink=0x%llx, Blink=0x%llx\n",
                (uint64_t)entry.InLoadOrderLinks.Flink, (uint64_t)entry.InLoadOrderLinks.Blink);

            // 移动到下一个条目
            // 根据链表类型选择正确的Flink字段
            if (listOffset == offsetof(X64_PEB_LDR_DATA, InLoadOrderModuleList))
                currentLink = (uint64_t)entry.InLoadOrderLinks.Flink;
            else if (listOffset == offsetof(X64_PEB_LDR_DATA, InMemoryOrderModuleList))
                currentLink = (uint64_t)entry.InMemoryOrderLinks.Flink;
            else if (listOffset == offsetof(X64_PEB_LDR_DATA, InInitializationOrderModuleList))
                currentLink = (uint64_t)entry.InInitializationOrderLinks.Flink;

            // 内存验证 - 防止无效指针
            if (currentLink == 0 || currentLink == entryAddr) {
                printf("错误：无效的Flink指针或自引用\n");
                break;
            }
        }

        if (entryCount == 0) {
            printf("警告：没有找到条目，但链表头指示非空\n");
        }
        else if (entryCount >= 100) {
            printf("警告：可能的无限循环，链表结构可能已损坏\n");
        }
        else {
            printf("\n总共找到 %d 个条目\n", entryCount);
        }
    }
}