fix up

ai_anti_malware/sandbox.h

@@ -111,6 +111,46 @@ class Sandbox {
     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;
@@ -167,12 +207,14 @@ class Sandbox {
         L"USERPROFILE=C:\\Users\\User",
         L"windir=C:\\Windows"};
     auto ResoveImport() -> void;
-    auto SetupVirtualMachine() -> void;
-    auto PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void;
-    auto processImportModule(const moudle_import* importModule) -> 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;         // 上次执行的区段索引
 };