#include "sandbox.h"


auto Sandbox::InitializeLdrData() -> void {
    if (m_peInfo->isX64 && m_peb64.Ldr == 0) {
        // 为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;

        // 初始化链表头 - 使用适当的类型转换
        LIST_ENTRY inLoadOrderList = {
            reinterpret_cast<LIST_ENTRY*>(
                ldrDataAddress +
                offsetof(X64_PEB_LDR_DATA, InLoadOrderModuleList)),
            reinterpret_cast<LIST_ENTRY*>(
                ldrDataAddress +
                offsetof(X64_PEB_LDR_DATA, InLoadOrderModuleList)) };
        ldrData.InLoadOrderModuleList = inLoadOrderList;

        LIST_ENTRY inMemoryOrderList = {
            reinterpret_cast<LIST_ENTRY*>(
                ldrDataAddress +
                offsetof(X64_PEB_LDR_DATA, InMemoryOrderModuleList)),
            reinterpret_cast<LIST_ENTRY*>(
                ldrDataAddress +
                offsetof(X64_PEB_LDR_DATA, InMemoryOrderModuleList)) };
        ldrData.InMemoryOrderModuleList = inMemoryOrderList;

        LIST_ENTRY inInitOrderList = {
            reinterpret_cast<LIST_ENTRY*>(
                ldrDataAddress +
                offsetof(X64_PEB_LDR_DATA, InInitializationOrderModuleList)),
            reinterpret_cast<LIST_ENTRY*>(
                ldrDataAddress +
                offsetof(X64_PEB_LDR_DATA, InInitializationOrderModuleList)) };
        ldrData.InInitializationOrderModuleList = inInitOrderList;

        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 = { 0 };
    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 };
    std::mbstowcs(nameBuffer, module->name, strlen(module->name));

    // 设置全路径
    entry.FullDllName.Length =
        static_cast<USHORT>(wcslen(nameBuffer) * sizeof(wchar_t));
    entry.FullDllName.MaximumLength = MAX_PATH * sizeof(wchar_t);
    entry.FullDllName.Buffer = reinterpret_cast<PWSTR>(fullNameAddress);

    // 设置基本名称
    entry.BaseDllName.Length =
        static_cast<USHORT>(wcslen(nameBuffer) * sizeof(wchar_t));
    entry.BaseDllName.MaximumLength = MAX_PATH * sizeof(wchar_t);
    entry.BaseDllName.Buffer = reinterpret_cast<PWSTR>(baseNameAddress);

    // 写入Unicode字符串
    uc_mem_write(m_ucEngine, fullNameAddress, nameBuffer,
        (wcslen(nameBuffer) + 1) * sizeof(wchar_t));
    uc_mem_write(m_ucEngine, baseNameAddress, nameBuffer,
        (wcslen(nameBuffer) + 1) * sizeof(wchar_t));

    return entry;
}

auto Sandbox::UpdateLdrLinks(const LDR_DATA_TABLE_ENTRY& entry,
    uint64_t entryAddress, X64_PEB_LDR_DATA& ldrData)
    -> void {
    // 更新LDR_DATA中的链表头
    ldrData.InLoadOrderModuleList.Flink = reinterpret_cast<LIST_ENTRY*>(
        entryAddress + offsetof(LDR_DATA_TABLE_ENTRY, InLoadOrderLinks));
    ldrData.InMemoryOrderModuleList.Flink = reinterpret_cast<LIST_ENTRY*>(
        entryAddress + offsetof(LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks));
    ldrData.InInitializationOrderModuleList.Flink =
        reinterpret_cast<LIST_ENTRY*>(
            entryAddress +
            offsetof(LDR_DATA_TABLE_ENTRY, InInitializationOrderLinks));

    // 写回更新后的LDR_DATA
    uc_mem_write(m_ucEngine, m_peb64.Ldr, &ldrData, sizeof(X64_PEB_LDR_DATA));
}

auto Sandbox::AddModuleToLdr(const std::shared_ptr<struct_moudle>& module)
-> void {
    if (!m_peInfo->isX64) {
        return;  // 暂时只处理64位
    }

    if (m_peb64.Ldr == 0) {
        InitializeLdrData();
    }

    // 为模块创建LDR_DATA_TABLE_ENTRY
    uint64_t entrySize = sizeof(LDR_DATA_TABLE_ENTRY) +
        MAX_PATH * 2;  // 额外空间用于Unicode字符串
    uint64_t entryAddress = m_pebEnd;
    m_pebEnd += entrySize;

    // 映射内存
    uc_mem_map(m_ucEngine, entryAddress, entrySize, UC_PROT_ALL);

    // 设置Unicode字符串地址
    uint64_t fullNameAddress = entryAddress + sizeof(LDR_DATA_TABLE_ENTRY);
    uint64_t baseNameAddress = fullNameAddress + MAX_PATH;

    // 创建并初始化LDR_DATA_TABLE_ENTRY
    auto entry =
        CreateLdrEntry(module, entryAddress, fullNameAddress, baseNameAddress);

    // 从PEB读取当前LDR_DATA结构
    X64_PEB_LDR_DATA ldrData;
    uc_mem_read(m_ucEngine, m_peb64.Ldr, &ldrData, sizeof(X64_PEB_LDR_DATA));

    // 设置链表指针
    entry.InLoadOrderLinks.Flink = reinterpret_cast<LIST_ENTRY*>(
        reinterpret_cast<uintptr_t>(ldrData.InLoadOrderModuleList.Flink));
    entry.InLoadOrderLinks.Blink = reinterpret_cast<LIST_ENTRY*>(
        m_peb64.Ldr + offsetof(X64_PEB_LDR_DATA, InLoadOrderModuleList));

    entry.InMemoryOrderLinks.Flink = reinterpret_cast<LIST_ENTRY*>(
        reinterpret_cast<uintptr_t>(ldrData.InMemoryOrderModuleList.Flink));
    entry.InMemoryOrderLinks.Blink = reinterpret_cast<LIST_ENTRY*>(
        m_peb64.Ldr + offsetof(X64_PEB_LDR_DATA, InMemoryOrderModuleList));

    entry.InInitializationOrderLinks.Flink =
        reinterpret_cast<LIST_ENTRY*>(reinterpret_cast<uintptr_t>(
            ldrData.InInitializationOrderModuleList.Flink));
    entry.InInitializationOrderLinks.Blink = reinterpret_cast<LIST_ENTRY*>(
        m_peb64.Ldr +
        offsetof(X64_PEB_LDR_DATA, InInitializationOrderModuleList));

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

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

    printf("Added module '%s' to LDR data tables at 0x%llx\n", module->name,
        entryAddress);
}