fix up

2025-03-09 00:06:37 +08:00
parent f7b0625bff
commit d2ed7936df
6 changed files with 615 additions and 36 deletions
--- a/ai_anti_malware/sandbox.cpp
+++ b/ai_anti_malware/sandbox.cpp
@@ -173,10 +173,30 @@ auto Sandbox::PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void {
            return;
        }
    }
+
+    // 创建新模块
+    auto newModule = CreateModuleInfo(dllName, moduleBase, moduleBase);
+
+    m_moduleList.push_back(newModule);
+    printf("push `%s` module to vm base: %llx vm size: %llx\n", newModule->name,
+           newModule->base, newModule->size);
+    uc_mem_map(m_ucEngine, newModule->base, newModule->size,
+               UC_PROT_READ | UC_PROT_EXEC);
+    uc_mem_write(m_ucEngine, newModule->base, (void*)moduleBase,
+                 newModule->size);
+    if (peconv::relocate_module((BYTE*)moduleBase, newModule->size,
+                                newModule->base) == false) {
+        throw std::runtime_error("Failed to relocate module");
+    }
+}
+
+auto Sandbox::CreateModuleInfo(const char* dllName, uint64_t moduleBase,
+                               uint64_t bufferAddress)
+    -> std::shared_ptr<struct_moudle> {
    // 解析PE头
-    auto* dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(moduleBase);
+    auto* dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(bufferAddress);
    auto* ntHeaders = reinterpret_cast<PIMAGE_NT_HEADERS>(
-        reinterpret_cast<LPBYTE>(moduleBase) + dosHeader->e_lfanew);
+        reinterpret_cast<LPBYTE>(bufferAddress) + dosHeader->e_lfanew);

    // 获取区段对齐值
    DWORD sectionAlignment =
@@ -191,7 +211,6 @@ auto Sandbox::PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void {
        sizeof(ntHeaders->FileHeader) +
        ntHeaders->FileHeader.SizeOfOptionalHeader);

-    // 创建新模块
    struct_moudle newModule{};
    strncpy(newModule.name, dllName, strlen(dllName));
    newModule.base =
@@ -216,9 +235,9 @@ auto Sandbox::PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void {
            protection |= UC_PROT_WRITE;

        // 计算区段大小
-        auto sectionSize =
-            AlignSize(max(section.Misc.VirtualSize, section.SizeOfRawData),
-                      sectionAlignment);
+        auto sectionSize = AlignToSectionAlignment(
+            max(section.Misc.VirtualSize, section.SizeOfRawData),
+            sectionAlignment);

        // 创建区段信息
        moudle_section newSection{};
@@ -233,17 +252,10 @@ auto Sandbox::PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void {
        std::cout << "[PE] " << dllName << " Section found: " << newSection.name
                  << '\n';
    }
-    m_moduleList.push_back(std::make_shared<struct_moudle>(newModule));
-    printf("push `%s` module to vm base: %llx vm size: %llx\n", newModule.name,
-           newModule.base, newModule.size);
-    uc_mem_map(m_ucEngine, newModule.base, newModule.size,
-               UC_PROT_READ | UC_PROT_EXEC);
-    uc_mem_write(m_ucEngine, newModule.base, (void*)moduleBase, newModule.size);
-    if (peconv::relocate_module((BYTE*)moduleBase, newModule.size,
-                                newModule.base) == false) {
-        throw std::runtime_error("Failed to relocate module");
-    }
+
+    return std::make_shared<struct_moudle>(newModule);
 }
+
 auto Sandbox::ResolveExport(uint64_t moduleBase)
    -> std::vector<std::shared_ptr<moudle_export>> {
    DWORD exportSize = 0;
@@ -295,6 +307,9 @@ auto Sandbox::ResolveExport(uint64_t moduleBase)
 }
 auto Sandbox::ResolveImportExports() -> void {
    for (auto module : m_moduleList) {
+        if (module->base == m_peInfo->RecImageBase) {
+            continue;
+        }
        const auto exports = ResolveExport(module->real_base);
        for (const auto item : exports) {
            printf("import export: [%s] %s => %llx\n", module->name, item->name,
@@ -388,8 +403,8 @@ auto Sandbox::SetupVirtualMachine() -> void {
    */
    uint64_t m_KSharedUserDataBase = 0x7FFE0000;
    uint64_t m_KSharedUserDataEnd = 0x7FFE0FFF;  // 0x7FFE2000
-    uint64_t m_KSharedUserDataSize =
-        AlignSize(m_KSharedUserDataEnd - m_KSharedUserDataBase, PAGE_SIZE);
+    uint64_t m_KSharedUserDataSize = AlignToSectionAlignment(
+        m_KSharedUserDataEnd - m_KSharedUserDataBase, PAGE_SIZE);

    uc_mem_map(m_ucEngine, m_KSharedUserDataBase, m_KSharedUserDataSize,
               UC_PROT_READ);
@@ -400,10 +415,10 @@ auto Sandbox::SetupVirtualMachine() -> void {
    m_pebBase = PEB_BASE;             // 进程PEB地址
    m_envBlockBase = ENV_BLOCK_BASE;  // 环境变量块地址
    // stack
-    m_stackBase =
-        AlignSize(this->m_peInfo->isX64 ? STACK_BASE_64 : STACK_BASE_32, 16);
-    m_stackSize =
-        AlignSize(this->m_peInfo->isX64 ? STACK_SIZE_64 : STACK_SIZE_32, 16);
+    m_stackBase = AlignToSectionAlignment(
+        this->m_peInfo->isX64 ? STACK_BASE_64 : STACK_BASE_32, 16);
+    m_stackSize = AlignToSectionAlignment(
+        this->m_peInfo->isX64 ? STACK_SIZE_64 : STACK_SIZE_32, 16);
    m_stackEnd = m_stackBase + m_stackSize;

    // heap
@@ -415,8 +430,10 @@ auto Sandbox::SetupVirtualMachine() -> void {
    if (this->m_peInfo->isX64) {
        // 设置64位PEB
        m_peb64.ImageBaseAddress = m_peInfo->RecImageBase;
-        m_pebEnd = m_pebBase + AlignSize(sizeof(X64PEB), PAGE_SIZE);
-        m_tebEnd = m_tebBase + AlignSize(sizeof(X64TEB), PAGE_SIZE);
+        m_pebEnd =
+            m_pebBase + AlignToSectionAlignment(sizeof(X64PEB), PAGE_SIZE);
+        m_tebEnd =
+            m_tebBase + AlignToSectionAlignment(sizeof(X64TEB), PAGE_SIZE);

        // 设置64位TEB
        m_teb64.ClientId.UniqueProcess = GetCurrentProcessId();
@@ -431,7 +448,8 @@ auto Sandbox::SetupVirtualMachine() -> void {
        // 设置GS基址结构
        m_gsBaseStruct.teb = m_tebBase;
        m_gsBaseStruct.peb = m_pebBase;
-        uint64_t gsAllocSize = AlignSize(sizeof(struct_gs_base), PAGE_SIZE);
+        uint64_t gsAllocSize =
+            AlignToSectionAlignment(sizeof(struct_gs_base), PAGE_SIZE);

        // 映射PEB到虚拟内存
        uc_mem_map(m_ucEngine, m_pebBase, m_pebEnd - m_pebBase,
@@ -456,8 +474,10 @@ auto Sandbox::SetupVirtualMachine() -> void {
    } else {
        // 设置32位PEB
        m_peb32.ImageBaseAddress = static_cast<ULONG>(m_peInfo->RecImageBase);
-        m_pebEnd = m_pebBase + AlignSize(sizeof(X32PEB), PAGE_SIZE);
-        m_tebEnd = m_tebBase + AlignSize(sizeof(X32TEB), PAGE_SIZE);
+        m_pebEnd =
+            m_pebBase + AlignToSectionAlignment(sizeof(X32PEB), PAGE_SIZE);
+        m_tebEnd =
+            m_tebBase + AlignToSectionAlignment(sizeof(X32TEB), PAGE_SIZE);

        // 设置32位TEB
        m_teb32.ClientId.UniqueProcess = GetCurrentProcessId();
@@ -491,7 +511,8 @@ auto Sandbox::SetupVirtualMachine() -> void {
        uc_reg_write(m_ucEngine, UC_X86_REG_MSR, &msr);
    }
    // 映射新的内存区域
-    size_t envSize = AlignSize(this->GetEnvStringsSize(), PAGE_SIZE);
+    size_t envSize =
+        AlignToSectionAlignment(this->GetEnvStringsSize(), PAGE_SIZE);
    printf("env block size: %llx\n", envSize);  // 添加调试输出
    uc_err envErr = uc_mem_map(m_ucEngine, m_envBlockBase, envSize,
                               UC_PROT_READ | UC_PROT_WRITE);
@@ -524,6 +545,13 @@ auto Sandbox::InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void {
        cs_close(&m_csHandle);  // 清理已分配的capstone资源
        throw std::runtime_error("Failed to initialize Unicorn");
    }
+    // 一定要确保他是第一个.
+    auto newModule =
+        CreateModuleInfo("huoji.exe", m_peInfo->RecImageBase,
+                         reinterpret_cast<uint64_t>(m_peInfo->peBuffer));
+    _ASSERTE(m_moduleList.size() == 0);
+    m_moduleList.push_back(newModule);
+
    ResoveImport();
    ResolveImportExports();

@@ -702,3 +730,342 @@ auto Sandbox::GetEnvString() -> std::vector<wchar_t> {
 auto Sandbox::GetEnvStringsSize() -> size_t {
    return GetEnvString().size() * sizeof(wchar_t);
 }
+
+auto Sandbox::getVirtualMemorySize(BYTE* peBuffer) -> size_t {
+    if (!peBuffer) {
+        return 0;
+    }
+
+    // 解析PE头
+    auto* dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(peBuffer);
+    if (dosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
+        return 0;
+    }
+    auto* ntHeaders = reinterpret_cast<PIMAGE_NT_HEADERS>(
+        reinterpret_cast<LPBYTE>(peBuffer) + dosHeader->e_lfanew);
+    // 获取区段头
+    auto* sectionHeader = reinterpret_cast<PIMAGE_SECTION_HEADER>(
+        reinterpret_cast<PUCHAR>(ntHeaders) + sizeof(ntHeaders->Signature) +
+        sizeof(ntHeaders->FileHeader) +
+        ntHeaders->FileHeader.SizeOfOptionalHeader);
+
+    DWORD minOffset = UINT_MAX;
+    DWORD totalSize = 0;
+
+    // 遍历所有区段
+    for (WORD i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++) {
+        const auto& section = sectionHeader[i];
+
+        // 查找最小虚拟地址偏移
+        if (section.VirtualAddress < minOffset) {
+            minOffset = section.VirtualAddress;
+        }
+
+        // 累加虚拟大小
+        totalSize += section.Misc.VirtualSize;
+    }
+
+    // 添加最小偏移到总大小
+    totalSize += minOffset;
+
+    return static_cast<size_t>(totalSize);
+}
+
+auto Sandbox::DumpPE() -> std::pair<std::unique_ptr<BYTE[]>, size_t> {
+    // 查找目标模块 - 这里我们使用主模块(通常是被分析的可执行文件)
+    std::shared_ptr<struct_moudle> targetModule = nullptr;
+    for (const auto& module : m_moduleList) {
+        if (strcmp(module->name, "huoji.exe") == 0) {
+            targetModule = module;
+            break;
+        }
+    }
+
+    if (!targetModule) {
+        throw std::runtime_error("No modules found to dump");
+    }
+
+    // 计算虚拟内存大小
+    auto virtualMemorySize = getVirtualMemorySize(m_peInfo->peBuffer);
+
+    // 创建用于存储转储数据的缓冲区
+    auto resultBuffer = std::make_unique<BYTE[]>(virtualMemorySize);
+
+    // 从虚拟机内存中读取PE文件
+    uc_err err = uc_mem_read(m_ucEngine, m_peInfo->RecImageBase,
+                             resultBuffer.get(), virtualMemorySize);
+    if (err != UC_ERR_OK) {
+        throw std::runtime_error("Failed to read memory during PE dump: " +
+                                 std::string(uc_strerror(err)));
+    }
+
+    // 确保PE头部的签名有效
+    auto* dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(resultBuffer.get());
+    if (dosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
+        throw std::runtime_error("Invalid DOS signature in dumped PE");
+    }
+
+    auto* ntHeaders = reinterpret_cast<PIMAGE_NT_HEADERS>(resultBuffer.get() +
+                                                          dosHeader->e_lfanew);
+    if (ntHeaders->Signature != IMAGE_NT_SIGNATURE) {
+        throw std::runtime_error("Invalid NT signature in dumped PE");
+    }
+
+    // 获取当前RIP/EIP作为新的入口点
+    uint64_t currentEntryPoint = 0;
+    if (this->GetCrossSectionExecution().size() > 0) {
+        currentEntryPoint = this->GetCrossSectionExecution()
+                                [this->GetCrossSectionExecution().size() - 1] -
+                            m_peInfo->RecImageBase;
+    }
+
+    PIMAGE_SECTION_HEADER sectionHeaders = nullptr;
+    WORD numberOfSections = 0;
+
+    // 处理32位或64位PE文件
+    if (m_peInfo->isX64) {
+        auto* optHeader64 =
+            &reinterpret_cast<PIMAGE_NT_HEADERS64>(ntHeaders)->OptionalHeader;
+        optHeader64->ImageBase = m_peInfo->RecImageBase;
+        if (currentEntryPoint != 0) {
+            // 修改入口点为当前执行位置
+            optHeader64->AddressOfEntryPoint =
+                static_cast<DWORD>(currentEntryPoint);
+        }
+
+        // 修改SizeOfImage
+        optHeader64->SizeOfImage = static_cast<DWORD>(AlignToSectionAlignment(
+            virtualMemorySize, optHeader64->SectionAlignment));
+
+        // 修改DllCharacteristics以移除ASLR标记
+        optHeader64->DllCharacteristics &=
+            ~IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE;
+
+        // 获取区段头信息
+        sectionHeaders = reinterpret_cast<PIMAGE_SECTION_HEADER>(
+            reinterpret_cast<ULONG_PTR>(ntHeaders) +
+            sizeof(ntHeaders->Signature) + sizeof(ntHeaders->FileHeader) +
+            ntHeaders->FileHeader.SizeOfOptionalHeader);
+        numberOfSections = ntHeaders->FileHeader.NumberOfSections;
+    } else {
+        auto* optHeader32 =
+            &reinterpret_cast<PIMAGE_NT_HEADERS32>(ntHeaders)->OptionalHeader;
+        optHeader32->ImageBase = static_cast<DWORD>(m_peInfo->RecImageBase);
+
+        if (currentEntryPoint != 0) {
+            // 修改入口点为当前执行位置
+            optHeader32->AddressOfEntryPoint =
+                static_cast<DWORD>(currentEntryPoint);
+        }
+
+        // 修改SizeOfImage
+        optHeader32->SizeOfImage = static_cast<DWORD>(AlignToSectionAlignment(
+            virtualMemorySize, optHeader32->SectionAlignment));
+
+        // 修改DllCharacteristics以移除ASLR标记
+        optHeader32->DllCharacteristics &=
+            ~IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE;
+
+        // 获取区段头信息
+        sectionHeaders = reinterpret_cast<PIMAGE_SECTION_HEADER>(
+            reinterpret_cast<ULONG_PTR>(ntHeaders) +
+            sizeof(ntHeaders->Signature) + sizeof(ntHeaders->FileHeader) +
+            ntHeaders->FileHeader.SizeOfOptionalHeader);
+        numberOfSections = ntHeaders->FileHeader.NumberOfSections;
+    }
+
+    // 更新代码基址和大小
+    UpdateBaseOfCode(sectionHeaders, ntHeaders, numberOfSections,
+                     static_cast<DWORD>(currentEntryPoint));
+
+    // 修复区段
+    FixSections(sectionHeaders, numberOfSections, virtualMemorySize);
+
+    // 创建一个ExportsMapper对象用于导入表修复
+    peconv::ExportsMapper exportsMap;
+
+    // 添加所有已加载模块到导出表映射中
+    for (const auto& module : m_moduleList) {
+        if (module->base == 0 || module->size == 0) {
+            continue;
+        }
+
+        // 创建临时缓冲区以存储模块内容
+        std::unique_ptr<BYTE[]> moduleBuffer =
+            std::make_unique<BYTE[]>(module->size);
+
+        // 从虚拟机内存读取模块内容
+        uc_err readErr = uc_mem_read(m_ucEngine, module->base,
+                                     moduleBuffer.get(), module->size);
+        if (readErr != UC_ERR_OK) {
+            printf(
+                "Warning: Could not read module %s for exports mapping: %s\n",
+                module->name, uc_strerror(readErr));
+            continue;
+        }
+
+        // 添加模块到导出表映射
+        exportsMap.add_to_lookup(module->name,
+                                 reinterpret_cast<HMODULE>(moduleBuffer.get()),
+                                 module->base);
+    }
+    //这里有一个严重的问题,就懒得处理了:
+    //壳里面吐出来的代码的导入表和壳的导入表不是同样一个.
+    //这个修的是壳的 导入表,所以导入表 修 不 全
+    //有个很简单的办法,需要搜索IAT结构,然后修改脱壳后的IAT的字段到壳的字段里面,然后再执行一次fix_imports
+    //懒得写了,家庭作业.自己完成
+    bool importsFixed = peconv::fix_imports(
+        resultBuffer.get(), virtualMemorySize, exportsMap, nullptr);
+    if (importsFixed) {
+        printf("PE file imports fixed successfully\n");
+    } else {
+        printf("Warning: Failed to fix PE file imports\n");
+    }
+
+    size_t out_size = 0;
+
+    // 重新计算校验和
+    if (m_peInfo->isX64) {
+        auto* optHeader64 =
+            &reinterpret_cast<PIMAGE_NT_HEADERS64>(ntHeaders)->OptionalHeader;
+        optHeader64->CheckSum =
+            CalculateChecksum(resultBuffer.get(), virtualMemorySize);
+    } else {
+        auto* optHeader32 =
+            &reinterpret_cast<PIMAGE_NT_HEADERS32>(ntHeaders)->OptionalHeader;
+        optHeader32->CheckSum =
+            CalculateChecksum(resultBuffer.get(), virtualMemorySize);
+    }
+
+    printf(
+        "PE file dumped successfully from address: 0x%llx, size: %zu bytes\n",
+        m_peInfo->RecImageBase, virtualMemorySize);
+    printf("Entry point set to: 0x%llx (RVA: 0x%llx)\n",
+           m_peInfo->RecImageBase + currentEntryPoint, currentEntryPoint);
+
+    return {std::move(resultBuffer), virtualMemorySize};
+}
+
+// 修复区段信息
+void Sandbox::FixSections(PIMAGE_SECTION_HEADER sectionHeaders,
+                          WORD numberOfSections, size_t virtualMemorySize) {
+    if (numberOfSections == 0 || sectionHeaders == nullptr) {
+        return;
+    }
+
+    // 修复每个区段的信息
+    for (WORD i = 0; i < numberOfSections - 1; i++) {
+        auto& currentSection = sectionHeaders[i];
+        auto& nextSection = sectionHeaders[i + 1];
+
+        // 修复大小，使之与下一个区段的起始地址对齐
+        currentSection.SizeOfRawData =
+            nextSection.VirtualAddress - currentSection.VirtualAddress;
+        currentSection.PointerToRawData = currentSection.VirtualAddress;
+        currentSection.Misc.VirtualSize = currentSection.SizeOfRawData;
+    }
+
+    // 修复最后一个区段
+    auto& lastSection = sectionHeaders[numberOfSections - 1];
+    lastSection.SizeOfRawData =
+        static_cast<DWORD>(virtualMemorySize) - lastSection.VirtualAddress;
+    lastSection.PointerToRawData = lastSection.VirtualAddress;
+    lastSection.Misc.VirtualSize = lastSection.SizeOfRawData;
+}
+
+// 计算校验和
+DWORD Sandbox::CalculateChecksum(const BYTE* peBuffer, size_t size) {
+    DWORD sum = 0;
+    const DWORD* ptr = reinterpret_cast<const DWORD*>(peBuffer);
+    const DWORD count = static_cast<DWORD>(size / sizeof(DWORD));
+
+    // 获取校验和字段的偏移
+    const auto dosHeader = (PIMAGE_DOS_HEADER)(peBuffer);
+    const auto ntHeaders = (PIMAGE_NT_HEADERS)(peBuffer + dosHeader->e_lfanew);
+    DWORD checksumOffset = dosHeader->e_lfanew +
+                           FIELD_OFFSET(IMAGE_NT_HEADERS, OptionalHeader) +
+                           FIELD_OFFSET(IMAGE_OPTIONAL_HEADER, CheckSum);
+
+    // 计算总和，跳过校验和字段本身
+    for (DWORD i = 0; i < count; i++) {
+        // 跳过校验和字段
+        if ((i * sizeof(DWORD)) == checksumOffset ||
+            (i * sizeof(DWORD)) == checksumOffset + sizeof(DWORD) - 1) {
+            continue;
+        }
+        sum += ptr[i];
+        // 处理溢出
+        if (sum < ptr[i]) {
+            sum++;
+        }
+    }
+
+    // 完成计算
+    sum = (sum & 0xFFFF) + (sum >> 16);
+    sum = (sum & 0xFFFF) + (sum >> 16);
+    sum = sum + static_cast<DWORD>(size);
+
+    return sum;
+}
+
+// 按区段对齐大小进行对齐
+DWORD Sandbox::AlignToSectionAlignment(size_t size, DWORD alignment) {
+    return static_cast<DWORD>(((size + alignment - 1) / alignment) * alignment);
+}
+
+// 更新代码基址和代码大小
+void Sandbox::UpdateBaseOfCode(PIMAGE_SECTION_HEADER sectionHeader,
+                               PIMAGE_NT_HEADERS ntHeaders,
+                               WORD numberOfSections, DWORD entryPoint) {
+    if (sectionHeader == nullptr || ntHeaders == nullptr ||
+        numberOfSections == 0) {
+        return;
+    }
+
+    DWORD baseOfCode = 0;
+    DWORD sizeOfCode = 0;
+    bool foundSection = false;
+
+    // 寻找包含入口点的区段
+    for (WORD i = 0; i < numberOfSections; i++) {
+        auto& section = sectionHeader[i];
+        if (entryPoint >= section.VirtualAddress &&
+            entryPoint < (section.VirtualAddress + section.Misc.VirtualSize)) {
+            baseOfCode = section.VirtualAddress;
+            sizeOfCode = section.Misc.VirtualSize;
+            foundSection = true;
+            break;
+        }
+    }
+
+    // 如果没有找到包含入口点的区段，使用第一个可执行区段
+    if (!foundSection) {
+        for (WORD i = 0; i < numberOfSections; i++) {
+            auto& section = sectionHeader[i];
+            if (section.Characteristics & IMAGE_SCN_MEM_EXECUTE) {
+                baseOfCode = section.VirtualAddress;
+                sizeOfCode = section.Misc.VirtualSize;
+                foundSection = true;
+                break;
+            }
+        }
+    }
+
+    // 更新NT头部信息
+    if (foundSection) {
+        if (ntHeaders->FileHeader.Machine == IMAGE_FILE_MACHINE_AMD64) {
+            // 64位PE
+            auto* optHeader64 =
+                &reinterpret_cast<PIMAGE_NT_HEADERS64>(ntHeaders)
+                     ->OptionalHeader;
+            optHeader64->BaseOfCode = baseOfCode;
+        } else {
+            // 32位PE
+            auto* optHeader32 =
+                &reinterpret_cast<PIMAGE_NT_HEADERS32>(ntHeaders)
+                     ->OptionalHeader;
+            optHeader32->BaseOfCode = baseOfCode;
+            optHeader32->SizeOfCode = sizeOfCode;
+        }
+    }
+}