update

ai_anti_malware/sandbox.h

@@ -1,6 +1,8 @@
 #pragma once
-#include "head.h"
+#include <functional>
+#include <map>
 
+#include "head.h"
 #define PAGE_SIZE 0x1000
 #define CF_MASK (1 << 0)
 #define PF_MASK (1 << 2)
@@ -17,8 +19,10 @@
 #define HEAP_ADDRESS_32 0x5000000
 #define HEAP_SIZE_32 0x5000000
 
-#define PEB_BASE 0x80000
+#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
@@ -45,20 +49,60 @@ enum class WinVer {
     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 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;
+    std::map<uint64_t, HeapSegment*> m_heapSegments;  // 堆段映射表
 
    private:
     std::shared_ptr<BasicPeInfo> m_peInfo;
@@ -75,19 +119,26 @@ class Sandbox {
     uint64_t m_heapSize;
     uint64_t m_heapEnd;
     uint64_t m_fakeBase;
-    struct_gs_base m_gsBaseStruct;
-    X64TEB m_teb64;
-    X64PEB m_peb64;
-    X32TEB m_teb32;
-    X32PEB m_peb32;
+    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;  // 存储命令行字符串
+
     auto ResoveImport() -> void;
     auto SetupVirtualMachine() -> void;
-    auto ResolveExport() -> void;
     auto PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void;
     auto processImportModule(const moudle_import* importModule) -> void;
+    auto ResolveImportExports() -> void;
+    auto ResolveExport(uint64_t moduleBase)
+        -> std::vector<std::shared_ptr<moudle_export>>;
+    auto InitApiHooks() -> void;
+    auto InitCommandLine(std::string commandLine) -> void;
 };