update

ai_anti_malware/sandbox.cpp

@@ -200,6 +200,7 @@ Sandbox::Sandbox() {
     m_ucEngine = nullptr;
     m_peInfo = nullptr;
     m_nextWfpEngineHandle = (HANDLE)0x1000;  // 初始化WFP引擎句柄
+    m_lastImpRead = { 0,0 };
 }
 
 Sandbox::~Sandbox() {
@@ -248,10 +249,13 @@ auto Sandbox::PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void {
             return;
         }
     }
-
+    if (m_usedModuleBase == 0) {
+        m_usedModuleBase = DLL_MODULE_BASE;
+    }
     // 创建新模块
-    auto newModule = CreateModuleInfo(dllName, moduleBase, moduleBase);
+    auto newModule = CreateModuleInfo(dllName, AlignSize(m_usedModuleBase, PAGE_SIZE), moduleBase, moduleBase);
 
+    m_usedModuleBase += PAGE_SIZE + newModule->size;
     m_moduleList.push_back(newModule);
     printf("push `%s` module to vm base: %llx vm size: %llx\n", newModule->name,
            newModule->base, newModule->size);
@@ -270,7 +274,7 @@ auto Sandbox::PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void {
     }
 }
 
-auto Sandbox::CreateModuleInfo(const char* dllName, uint64_t moduleBase,
+auto Sandbox::CreateModuleInfo(const char* dllName, uint64_t moduleBase, uint64_t realModuleBase,
                                uint64_t bufferAddress)
     -> std::shared_ptr<struct_moudle> {
     // 解析PE头
@@ -293,9 +297,8 @@ auto Sandbox::CreateModuleInfo(const char* dllName, uint64_t moduleBase,
 
     struct_moudle newModule{};
     strncpy(newModule.name, dllName, strlen(dllName));
-    newModule.base =
-        this->m_peInfo->isX64 ? moduleBase : static_cast<uint32_t>(moduleBase);
-    newModule.real_base = moduleBase;
+    newModule.base = moduleBase;
+    newModule.real_base = realModuleBase;
     newModule.entry = ntHeaders->OptionalHeader.AddressOfEntryPoint;
     newModule.size = ntHeaders->OptionalHeader.SizeOfImage;
     // 处理区段
@@ -633,7 +636,7 @@ auto Sandbox::InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void {
     }
     // 一定要确保他是第一个.
     auto newModule =
-        CreateModuleInfo("huoji.exe", m_peInfo->RecImageBase,
+        CreateModuleInfo("huoji.exe", m_peInfo->RecImageBase, m_peInfo->RecImageBase,
                          reinterpret_cast<uint64_t>(m_peInfo->peBuffer));
     _ASSERTE(m_moduleList.size() == 0);
     m_moduleList.push_back(newModule);