awesome_anti_virus_engine/ai_anti_malware/sandbox_dump_pe.cpp

#include "sandbox.h"

auto Sandbox::DumpPE() -> std::pair<std::unique_ptr<BYTE[]>, size_t> {
    // <20><><EFBFBD><EFBFBD>Ŀ<EFBFBD><C4BF>ģ<EFBFBD><C4A3> - <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ʹ<EFBFBD><CAB9><EFBFBD><EFBFBD>ģ<EFBFBD><C4A3>(ͨ<><CDA8><EFBFBD>DZ<EFBFBD><C7B1><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ŀ<EFBFBD>ִ<EFBFBD><D6B4><EFBFBD>ļ<EFBFBD>)
    std::shared_ptr<struct_moudle> targetModule = nullptr;
    for (const auto& module : m_moduleList) {
        if ((*module).name == "HUOJI.EXE") {
            targetModule = module;
            break;
        }
    }

    if (!targetModule) {
        throw std::runtime_error("No modules found to dump");
    }

    // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ڴ<EFBFBD><DAB4><EFBFBD>С
    auto virtualMemorySize = getVirtualMemorySize(m_peInfo->peBuffer);

    // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ڴ洢ת<E6B4A2><D7AA><EFBFBD><EFBFBD><EFBFBD>ݵĻ<DDB5><C4BB><EFBFBD><EFBFBD><EFBFBD>
    auto resultBuffer = std::make_unique<BYTE[]>(virtualMemorySize);

    // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ڴ<EFBFBD><DAB4>ж<EFBFBD>ȡPE<50>ļ<EFBFBD>
    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)));
    }

    // ȷ<><C8B7>PEͷ<45><CDB7><EFBFBD><EFBFBD>ǩ<EFBFBD><C7A9><EFBFBD><EFBFBD>Ч
    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");
    }

    // <20><>ȡ<EFBFBD><C8A1>ǰRIP/EIP<49><50>Ϊ<EFBFBD>µ<EFBFBD><C2B5><EFBFBD><EFBFBD>ڵ<EFBFBD>
    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;

    // <20><><EFBFBD><EFBFBD>32λ<32><CEBB>64λPE<50>ļ<EFBFBD>
    if (m_peInfo->isX64) {
        auto* optHeader64 =
            &reinterpret_cast<PIMAGE_NT_HEADERS64>(ntHeaders)->OptionalHeader;
        optHeader64->ImageBase = m_peInfo->RecImageBase;
        if (currentEntryPoint != 0) {
            // <20>޸<EFBFBD><DEB8><EFBFBD><EFBFBD>ڵ<EFBFBD>Ϊ<EFBFBD><CEAA>ǰִ<C7B0><D6B4>λ<EFBFBD><CEBB>
            optHeader64->AddressOfEntryPoint =
                static_cast<DWORD>(currentEntryPoint);
        }

        // <20>޸<EFBFBD>SizeOfImage
        optHeader64->SizeOfImage = static_cast<DWORD>(AlignToSectionAlignment(
            virtualMemorySize, optHeader64->SectionAlignment));

        // <20>޸<EFBFBD>DllCharacteristics<63><73><EFBFBD>Ƴ<EFBFBD>ASLR<4C><52><EFBFBD><EFBFBD>
        optHeader64->DllCharacteristics &=
            ~IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE;

        // <20><>ȡ<EFBFBD><C8A1><EFBFBD><EFBFBD>ͷ<EFBFBD><CDB7>Ϣ
        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) {
            // <20>޸<EFBFBD><DEB8><EFBFBD><EFBFBD>ڵ<EFBFBD>Ϊ<EFBFBD><CEAA>ǰִ<C7B0><D6B4>λ<EFBFBD><CEBB>
            optHeader32->AddressOfEntryPoint =
                static_cast<DWORD>(currentEntryPoint);
        }

        // <20>޸<EFBFBD>SizeOfImage
        optHeader32->SizeOfImage = static_cast<DWORD>(AlignToSectionAlignment(
            virtualMemorySize, optHeader32->SectionAlignment));

        // <20>޸<EFBFBD>DllCharacteristics<63><73><EFBFBD>Ƴ<EFBFBD>ASLR<4C><52><EFBFBD><EFBFBD>
        optHeader32->DllCharacteristics &=
            ~IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE;

        // <20><>ȡ<EFBFBD><C8A1><EFBFBD><EFBFBD>ͷ<EFBFBD><CDB7>Ϣ
        sectionHeaders = reinterpret_cast<PIMAGE_SECTION_HEADER>(
            reinterpret_cast<ULONG_PTR>(ntHeaders) +
            sizeof(ntHeaders->Signature) + sizeof(ntHeaders->FileHeader) +
            ntHeaders->FileHeader.SizeOfOptionalHeader);
        numberOfSections = ntHeaders->FileHeader.NumberOfSections;
    }

    // <20><><EFBFBD>´<EFBFBD><C2B4><EFBFBD><EFBFBD><EFBFBD>ַ<EFBFBD>ʹ<EFBFBD>С
    UpdateBaseOfCode(sectionHeaders, ntHeaders, numberOfSections,
        static_cast<DWORD>(currentEntryPoint));

    // <20>޸<EFBFBD><DEB8><EFBFBD><EFBFBD><EFBFBD>
    FixSections(sectionHeaders, numberOfSections, virtualMemorySize);

    // <20><><EFBFBD><EFBFBD>һ<EFBFBD><D2BB>ExportsMapper<65><72><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ڵ<EFBFBD><DAB5><EFBFBD><EFBFBD><EFBFBD><EFBFBD>޸<EFBFBD>
    peconv::ExportsMapper exportsMap;

    // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ѽ<EFBFBD><D1BC><EFBFBD>ģ<EFBFBD>鵽<EFBFBD><E9B5BD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ӳ<EFBFBD><D3B3><EFBFBD><EFBFBD>
    for (const auto& module : m_moduleList) {
        if (module->base == 0 || module->size == 0) {
            continue;
        }

        // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ʱ<EFBFBD><CAB1><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Դ洢ģ<E6B4A2><C4A3><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
        std::unique_ptr<BYTE[]> moduleBuffer =
            std::make_unique<BYTE[]>(module->size);

        // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ڴ<EFBFBD><DAB4><EFBFBD>ȡģ<C8A1><C4A3><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
        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;
        }

        // <20><><EFBFBD><EFBFBD>ģ<EFBFBD>鵽<EFBFBD><E9B5BD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ӳ<EFBFBD><D3B3>
        exportsMap.add_to_lookup(module->name,
            reinterpret_cast<HMODULE>(moduleBuffer.get()),
            module->base);
    }
    // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>һ<EFBFBD><D2BB><EFBFBD><EFBFBD><EFBFBD>ص<EFBFBD><D8B5><EFBFBD><EFBFBD><EFBFBD>,<2C><><EFBFBD><EFBFBD><EFBFBD>ô<EFBFBD><C3B4><EFBFBD><EFBFBD><EFBFBD>:
    // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>³<EFBFBD><C2B3><EFBFBD><EFBFBD>Ĵ<EFBFBD><C4B4><EFBFBD><EFBFBD>ĵ<EFBFBD><C4B5><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ϳǵĵ<C7B5><C4B5><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ͬ<EFBFBD><CDAC>һ<EFBFBD><D2BB>.
    // <20><><EFBFBD><EFBFBD><EFBFBD>޵<EFBFBD><DEB5>ǿǵ<C7BF> <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>,<2C><><EFBFBD>Ե<EFBFBD><D4B5><EFBFBD><EFBFBD><EFBFBD> <20><> <20><> ȫ
    // <20>и<EFBFBD><D0B8>ܼ򵥵İ취,<2C><>Ҫ<EFBFBD><D2AA><EFBFBD><EFBFBD>IAT<41>ṹ,Ȼ<><C8BB><EFBFBD>޸<EFBFBD><DEB8>ѿǺ<D1BF><C7BA><EFBFBD>IAT<41><54><EFBFBD>ֶε<D6B6><CEB5>ǵ<EFBFBD><C7B5>ֶ<EFBFBD><D6B6><EFBFBD><EFBFBD><EFBFBD>,Ȼ<><C8BB><EFBFBD><EFBFBD>ִ<EFBFBD><D6B4>һ<EFBFBD><D2BB>fix_imports
    // <20><><EFBFBD><EFBFBD>д<EFBFBD><D0B4>,<2C><>ͥ<EFBFBD><CDA5>ҵ.<2E>Լ<EFBFBD><D4BC><EFBFBD><EFBFBD><EFBFBD>
    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;

    // <20><><EFBFBD>¼<EFBFBD><C2BC><EFBFBD>У<EFBFBD><D0A3><EFBFBD><EFBFBD>
    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 };
}

// <20>޸<EFBFBD><DEB8><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ϣ
void Sandbox::FixSections(PIMAGE_SECTION_HEADER sectionHeaders,
    WORD numberOfSections, size_t virtualMemorySize) {
    if (numberOfSections == 0 || sectionHeaders == nullptr) {
        return;
    }

    // <20>޸<EFBFBD>ÿ<EFBFBD><C3BF><EFBFBD><EFBFBD><EFBFBD>ε<EFBFBD><CEB5><EFBFBD>Ϣ
    for (WORD i = 0; i < numberOfSections - 1; i++) {
        auto& currentSection = sectionHeaders[i];
        auto& nextSection = sectionHeaders[i + 1];

        // <20>޸<EFBFBD><DEB8><EFBFBD>С<EFBFBD><D0A1>ʹ֮<CAB9><D6AE><EFBFBD><EFBFBD>һ<EFBFBD><D2BB><EFBFBD><EFBFBD><EFBFBD>ε<EFBFBD><CEB5><EFBFBD>ʼ<EFBFBD><CABC>ַ<EFBFBD><D6B7><EFBFBD><EFBFBD>
        currentSection.SizeOfRawData =
            nextSection.VirtualAddress - currentSection.VirtualAddress;
        currentSection.PointerToRawData = currentSection.VirtualAddress;
        currentSection.Misc.VirtualSize = currentSection.SizeOfRawData;
    }

    // <20>޸<EFBFBD><DEB8><EFBFBD><EFBFBD><EFBFBD>һ<EFBFBD><D2BB><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
    auto& lastSection = sectionHeaders[numberOfSections - 1];
    lastSection.SizeOfRawData =
        static_cast<DWORD>(virtualMemorySize) - lastSection.VirtualAddress;
    lastSection.PointerToRawData = lastSection.VirtualAddress;
    lastSection.Misc.VirtualSize = lastSection.SizeOfRawData;
}

// <20><><EFBFBD><EFBFBD>У<EFBFBD><D0A3><EFBFBD><EFBFBD>
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));

    // <20><>ȡУ<C8A1><D0A3><EFBFBD><EFBFBD><EFBFBD>ֶε<D6B6>ƫ<EFBFBD><C6AB>
    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);

    // <20><><EFBFBD><EFBFBD><EFBFBD>ܺͣ<DCBA><CDA3><EFBFBD><EFBFBD><EFBFBD>У<EFBFBD><D0A3><EFBFBD><EFBFBD><EFBFBD>ֶα<D6B6><CEB1><EFBFBD>
    for (DWORD i = 0; i < count; i++) {
        // <20><><EFBFBD><EFBFBD>У<EFBFBD><D0A3><EFBFBD><EFBFBD><EFBFBD>ֶ<EFBFBD>
        if ((i * sizeof(DWORD)) == checksumOffset ||
            (i * sizeof(DWORD)) == checksumOffset + sizeof(DWORD) - 1) {
            continue;
        }
        sum += ptr[i];
        // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
        if (sum < ptr[i]) {
            sum++;
        }
    }

    // <20><><EFBFBD>ɼ<EFBFBD><C9BC><EFBFBD>
    sum = (sum & 0xFFFF) + (sum >> 16);
    sum = (sum & 0xFFFF) + (sum >> 16);
    sum = sum + static_cast<DWORD>(size);

    return sum;
}

// <20><><EFBFBD><EFBFBD><EFBFBD>ζ<EFBFBD><CEB6><EFBFBD><EFBFBD><EFBFBD>С<EFBFBD><D0A1><EFBFBD>ж<EFBFBD><D0B6><EFBFBD>
DWORD Sandbox::AlignToSectionAlignment(size_t size, DWORD alignment) {
    return static_cast<DWORD>(((size + alignment - 1) / alignment) * alignment);
}

// <20><><EFBFBD>´<EFBFBD><C2B4><EFBFBD><EFBFBD><EFBFBD>ַ<EFBFBD>ʹ<EFBFBD><CDB4><EFBFBD><EFBFBD><EFBFBD>С
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;

    // Ѱ<>Ұ<EFBFBD><D2B0><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ڵ<EFBFBD><DAB5><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
    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;
        }
    }

    // <20><><EFBFBD><EFBFBD>û<EFBFBD><C3BB><EFBFBD>ҵ<EFBFBD><D2B5><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ڵ<EFBFBD><DAB5><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Σ<EFBFBD>ʹ<EFBFBD>õ<EFBFBD>һ<EFBFBD><D2BB><EFBFBD><EFBFBD>ִ<EFBFBD><D6B4><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
    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;
            }
        }
    }

    // <20><><EFBFBD><EFBFBD>NTͷ<54><CDB7><EFBFBD><EFBFBD>Ϣ
    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;
        }
    }
}