awesome_anti_virus_engine/ai_anti_malware/sandbox.h

#pragma once
#include <functional>
#include <map>

#include "head.h"
#define PAGE_SIZE 0x1000
#define CF_MASK (1 << 0)
#define PF_MASK (1 << 2)
#define ZF_MASK (1 << 6)
#define SF_MASK (1 << 7)
#define OF_MASK (1 << 11)
#define ALL_MASK (OF_MASK | SF_MASK | ZF_MASK | PF_MASK | CF_MASK)
// 随便瞎JB写的
#define STACK_BASE_64 0x14A0000
#define STACK_BASE_32 0x14A0000
#define STACK_SIZE_64 0x40000
#define STACK_SIZE_32 0x40000
#define HEAP_ADDRESS_64 0x500000000
#define HEAP_SIZE_64 0x5000000
#define HEAP_ADDRESS_32 0x5000000
#define HEAP_SIZE_32 0x5000000
#define ENV_BLOCK_BASE 0x50000

#define PEB_BASE 0x90000
#define TEB_BASE 0x90000
#define CMDLINE_ADDRESS 0x100000   // 命令行字符串的固定地址
#define CMDLINEW_ADDRESS 0x110000  // 宽字符命令行字符串的固定地址

#define X86_GDT_ADDR 0x30000
#define X86_GDT_LIMIT 0x1000
#define X86_GDT_ENTRY_SIZE 0x8

#define API_FUNCTION_SIZE 8
#define PAGE_ALIGN(Va) (ULONG_PTR)(Va) & ~(PAGE_SIZE - 1)
#define PAGE_ALIGN_64(Va) (Va) & ~(0x1000ull - 1)
#define PAGE_ALIGN_64k(Va) ((Va)) & ~(0x10000ull - 1)
#define AlignSize(Size, Align) (Size + Align - 1) / Align* Align

enum class WinVer {
    kWin7 = 0x0610,
    kWin7SP1 = 0x0611,
    kWin8 = 0x0620,
    kWin81 = 0x0630,
    kWin10 = 0x0A00,
    kWin10RS1 = 0x0A01,   // Anniversary update
    kWin10RS2 = 0x0A02,   // Creators update
    kWin10RS3 = 0x0A03,   // Fall creators update
    kWin10RS4 = 0x0A04,   // Spring creators update
    kWin10RS5 = 0x0A05,   // October 2018 update
    kWin1019H1 = 0x0A06,  // May 2019 update 19H1
    kWin1019H2 = 0x0A07,  // November 2019 update 19H2
    kWin1020H1 = 0x0A08   // April 2020 update 20H1
};
struct _fakeApi {
    std::function<void(void*, uc_engine*, uint64_t)> func;
    uint32_t paramCount;
};

// 添加堆管理相关的结构定义
struct HeapBlock {
    uint64_t address;  // 块的起始地址
    size_t size;       // 块的大小
    bool is_free;      // 是否是空闲块
    HeapBlock* next;   // 下一个块
    HeapBlock* prev;   // 上一个块
};

struct HeapSegment {
    uint64_t base;      // 堆段的基址
    size_t size;        // 堆段的总大小
    HeapBlock* blocks;  // 块链表
};

class Sandbox {
    friend class cFixImprot;  // 声明cFixImprot为友元类
   public:
    Sandbox();
    ~Sandbox();

    // Public methods
    auto InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void;
    auto Run() -> void;
    auto GetCapstoneHandle() const -> csh { return m_csHandle; }
    auto GetUnicornHandle() const -> uc_engine* { return m_ucEngine; }
    auto GetPeInfo() const -> std::shared_ptr<BasicPeInfo> { return m_peInfo; }
    auto GetModuleList() const -> std::vector<std::shared_ptr<struct_moudle>> {
        return m_moduleList;
    }
    auto EmulateApi(uc_engine* uc, uint64_t address, uint64_t rip,
                    std::string ApiName) -> void;
    auto GetPeb32() -> X32PEB* { return &m_peb32; }
    auto GetPeb64() -> X64PEB* { return &m_peb64; }
    auto GetTeb32() -> X32TEB* { return &m_teb32; }
    auto GetTeb64() -> X64TEB* { return &m_teb64; }
    auto GetCommandLine() const -> const char* { return m_commandLine.c_str(); }
    auto GetCommandLineAddress() const -> uint64_t { return CMDLINE_ADDRESS; }
    auto GetCommandLineWAddress() const -> uint64_t { return CMDLINEW_ADDRESS; }
    auto GetEnvStrings() const -> std::vector<std::wstring> {
        return envStrings;
    }
    auto GetEnvString() -> std::vector<wchar_t>;
    auto GetEnvStringsSize() -> size_t;
    auto InitCommandLine() -> void;

    // 堆管理相关的公共方法
    auto CreateHeapSegment(uint64_t base, size_t size) -> HeapSegment*;
    auto AllocateFromSegment(HeapSegment* segment, size_t size) -> uint64_t;
    auto FreeBlock(uint64_t address) -> bool;
    auto FindHeapSegment(uint64_t address) -> HeapSegment*;
    auto MergeBlocks(HeapBlock* block) -> void;
    auto SplitBlock(HeapBlock* block, size_t size) -> void;
    auto GetEnvBlockBase() const -> uint64_t { return m_envBlockBase; }
    std::map<uint64_t, HeapSegment*> m_heapSegments;  // 堆段映射表
    auto GetHeapBlocks() const -> std::map<uint64_t, HeapSegment*> {
        return m_heapSegments;
    }

    // 从内存中提取PE文件并修复重定位和导入表，返回原始PE的缓冲区
    auto DumpPE() -> std::pair<std::unique_ptr<BYTE[]>, size_t>;

    // 计算PE文件的虚拟内存大小
    auto getVirtualMemorySize(BYTE* peBuffer) -> size_t;

    // 修复PE区段信息
    void FixSections(PIMAGE_SECTION_HEADER sectionHeader, WORD numberOfSections,
                     size_t virtualMemorySize);

    // 更新代码基址和大小
    void UpdateBaseOfCode(PIMAGE_SECTION_HEADER sectionHeader,
                          PIMAGE_NT_HEADERS ntHeaders, WORD numberOfSections,
                          DWORD entryPoint);

    // 对齐到区段对齐值
    DWORD AlignToSectionAlignment(size_t size, DWORD alignment);

    // 计算PE校验和
    DWORD CalculateChecksum(const BYTE* buffer, size_t size);

    auto SetupVirtualMachine() -> void;
    auto PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void;
    auto processImportModule(const moudle_import* importModule) -> void;
    auto GetCrossSectionExecution() -> std::vector<uint64_t> {
        return m_crossSectionExecution;
    }
    auto GetLastExecuteSectionIndex() -> uint64_t {
        return m_lastExecuteSectionIndex;
    }
    auto SetLastExecuteSectionIndex(uint64_t index) -> void {
        m_lastExecuteSectionIndex = index;
    }
    auto SetCrossSectionExecution(uint64_t address) -> void {
        return m_crossSectionExecution.push_back(address);
    }

   private:
    std::shared_ptr<BasicPeInfo> m_peInfo;
    uint64_t m_gsBase;
    uint64_t m_pebBase;
    uint64_t m_pebEnd;
    uint64_t m_tebBase;
    uint64_t m_tebEnd;
    PVOID m_stackBuffer;  // 没有释放
    uint64_t m_stackBase;
    uint64_t m_stackSize;
    uint64_t m_stackEnd;
    uint64_t m_heapBase;
    uint64_t m_heapSize;
    uint64_t m_heapEnd;
    uint64_t m_fakeBase;
    uint64_t m_envBlockBase;
    struct_gs_base m_gsBaseStruct = {0};
    X64TEB m_teb64 = {0};
    X64PEB m_peb64 = {0};
    X32TEB m_teb32 = {0};
    X32PEB m_peb32 = {0};
    csh m_csHandle;         // Capstone handle
    uc_engine* m_ucEngine;  // Unicorn engine handle
    std::vector<std::shared_ptr<moudle_import>> m_impFuncDict;
    std::vector<std::shared_ptr<moudle_export>> m_exportFuncDict;
    std::vector<std::shared_ptr<struct_moudle>> m_moduleList;
    std::map<std::string, std::shared_ptr<_fakeApi>> api_map;
    std::string m_commandLine;  // 存储命令行字符串
    // 创建一些基本的环境变量
    std::vector<std::wstring> envStrings = {
        L"ALLUSERSPROFILE=C:\\ProgramData",
        L"APPDATA=C:\\Users\\User\\AppData\\Roaming",
        L"CommonProgramFiles=C:\\Program Files\\Common Files",
        L"COMPUTERNAME=DESKTOP",
        L"ComSpec=C:\\Windows\\system32\\cmd.exe",
        L"HOMEDRIVE=C:",
        L"HOMEPATH=\\Users\\User",
        L"LOCALAPPDATA=C:\\Users\\User\\AppData\\Local",
        L"NUMBER_OF_PROCESSORS=8",
        L"OS=Windows_NT",
        L"Path=C:\\Windows\\system32;C:\\Windows;C:\\Windows\\System32\\Wbem",
        L"PATHEXT=.COM;.EXE;.BAT;.CMD;.VBS;.VBE;.JS;.JSE;.WSF;.WSH;.MSC",
        L"PROCESSOR_ARCHITECTURE=AMD64",
        L"ProgramData=C:\\ProgramData",
        L"ProgramFiles=C:\\Program Files",
        L"PROMPT=$P$G",
        L"SystemDrive=C:",
        L"SystemRoot=C:\\Windows",
        L"TEMP=C:\\Users\\User\\AppData\\Local\\Temp",
        L"TMP=C:\\Users\\User\\AppData\\Local\\Temp",
        L"USERDOMAIN=DESKTOP",
        L"USERNAME=User",
        L"USERPROFILE=C:\\Users\\User",
        L"windir=C:\\Windows"};
    auto ResoveImport() -> void;
    auto ResolveImportExports() -> void;
    auto CreateModuleInfo(const char* dllName, uint64_t moduleBase,
                          uint64_t bufferAddress)
        -> std::shared_ptr<struct_moudle>;
    auto ResolveExport(uint64_t moduleBase)
        -> std::vector<std::shared_ptr<moudle_export>>;
    auto InitApiHooks() -> void;
    auto InitCommandLine(std::string commandLine) -> void;
    std::vector<uint64_t> m_crossSectionExecution;  // 记录跨区段执行地址
    uint64_t m_lastExecuteSectionIndex = 0;         // 上次执行的区段索引
    uint64_t m_KSharedUserDataBase{0};
    uint64_t m_KSharedUserDataSize{0};
};