update

.vscode/settings.json

@@ -71,6 +71,7 @@
         "stop_token": "cpp",
         "thread": "cpp",
         "xlocbuf": "cpp",
-        "xlocmes": "cpp"
+        "xlocmes": "cpp",
+        "unordered_set": "cpp"
     }
 }

ai_anti_malware/ai_anti_malware.cpp

@@ -314,7 +314,7 @@ auto doMalwareScan(int argc, char* argv[]) -> void {
 
 int doSandbox(int argc, char* argv[]) {
 
-    std::string filePath = "C:\\opengl32.dll";
+    std::string filePath = "Z:\\opengl32.dll";
     auto peInfo = getPeInfo(filePath);
     if (peInfo == nullptr) {
         std::cout << "无法加载PE文件: " << filePath << std::endl;

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 =
+    newModule.base = moduleBase;
-        this->m_peInfo->isX64 ? moduleBase : static_cast<uint32_t>(moduleBase);
+    newModule.real_base = realModuleBase;
-    newModule.real_base = moduleBase;
     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);

ai_anti_malware/sandbox.h

@@ -19,6 +19,7 @@
 #define HEAP_ADDRESS_32 0x5000000
 #define HEAP_SIZE_32 0x5000000
 #define ENV_BLOCK_BASE 0x50000
+#define DLL_MODULE_BASE 0x130000
 
 #define PEB_BASE 0x90000
 #define TEB_BASE 0x90000
@@ -96,7 +97,6 @@ class Sandbox {
     ~Sandbox();
     std::map<uint64_t, size_t>
         process_enum_state;  // 用于跟踪每个句柄的枚举状态
-
     // Public methods
     auto InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void;
     auto Run(uint64_t address = 0) -> void;
@@ -107,7 +107,7 @@ class Sandbox {
         return m_moduleList;
     }
     auto EmulateApi(uc_engine* uc, uint64_t address, uint64_t rip,
-                    std::string ApiName) -> void;
+                    std::string ApiName) -> bool;
     auto GetPeb32() -> X32PEB* { return &m_peb32; }
     auto GetPeb64() -> X64PEB* { return &m_peb64; }
     auto GetTeb32() -> X32TEB* { return &m_teb32; }
@@ -232,9 +232,20 @@ class Sandbox {
         m_nextWfpEngineHandle = (HANDLE)((uint64_t)m_nextWfpEngineHandle + 1);
         return handle;
     }
+    auto GetImpFuncDict() -> std::vector<std::shared_ptr<moudle_import>> {
+        return m_impFuncDict;
+    }
+    auto GetLastImpRead() -> std::pair<uint64_t, std::shared_ptr<moudle_import>> {
+        return m_lastImpRead;
+    }
+    auto SetLastImpRead(uint64_t address, std::shared_ptr<moudle_import> imp) -> void {
+        m_lastImpRead = { address, imp };
+    }
 
    private:
     std::shared_ptr<BasicPeInfo> m_peInfo;
+    std::pair<uint64_t, std::shared_ptr<moudle_import>> m_lastImpRead;
+    uint64_t m_usedModuleBase;
     uint64_t m_gsBase;
     uint64_t m_pebBase;
     uint64_t m_pebEnd;
@@ -289,7 +300,7 @@ class Sandbox {
         L"windir=C:\\Windows"};
     auto ResoveImport() -> void;
     auto ResolveImportExports() -> void;
-    auto CreateModuleInfo(const char* dllName, uint64_t moduleBase,
+    auto CreateModuleInfo(const char* dllName, uint64_t moduleBase, uint64_t realModuleBase,
                           uint64_t bufferAddress)
         -> std::shared_ptr<struct_moudle>;
     auto ResolveExport(uint64_t moduleBase)

ai_anti_malware/sandbox_api_emu.cpp

@@ -1497,7 +1497,7 @@ auto Sandbox::InitApiHooks() -> void {
     };
 }
 auto Sandbox::EmulateApi(uc_engine* uc, uint64_t address, uint64_t rip,
-                         std::string ApiName) -> void {
+                         std::string ApiName) -> bool {
     auto it = api_map.find(ApiName);
     if (it != api_map.end()) {
         it->second->func(this, uc, address);
@@ -1545,9 +1545,9 @@ auto Sandbox::EmulateApi(uc_engine* uc, uint64_t address, uint64_t rip,
             uc_reg_write(uc, UC_X86_REG_ESP, &esp);
         }
 
-        return;
+        return true;
     }
     printf("ApiName: %s not found\n", ApiName.c_str());
-    uc_emu_stop(uc);
+    //uc_emu_stop(uc);
-    return;
+    return false;
 }

ai_anti_malware/sandbox_api_stl.cpp

@@ -2087,10 +2087,6 @@ auto Api_RtlFormatCurrentUserKeyPath(void* sandbox, uc_engine* uc,
     wchar_t userKeyPath[256] =
         L"\\Registry\\User\\S-1-5-21-1234567890-1234567890-1234567890-1001";
 
-    // 模拟查询MCP服务器，这里可以根据需要实现实际的服务器查询
-    // 例如可以修改userKeyPath或在日志中记录查询操作
-    printf("[*] RtlFormatCurrentUserKeyPath: 正在查询MCP服务器获取用户SID\n");
-
     // 分配内存用于存储路径字符串
     size_t pathLen = wcslen(userKeyPath);
     size_t bufferSize = (pathLen + 1) * sizeof(wchar_t);

ai_anti_malware/sandbox_callbacks.cpp

@@ -56,16 +56,26 @@ void handleCodeRun(uc_engine* uc, uint64_t address, uint32_t size,
     if (currentSectionIndex >= 0) {
         sandbox->SetLastExecuteSectionIndex(currentSectionIndex);
     }
+    auto [lastReadImpAddr, lastImp] = sandbox->GetLastImpRead();
+    if (lastImp != nullptr && currentRip == lastReadImpAddr) {
+        printf(
+            "direct call function [%s]%s at file address: %llx lastRip: "
+            "%llx\n",
+            lastImp->dll_name,
+            lastImp->name, address, lastRip);
+        sandbox->EmulateApi(uc, lastReadImpAddr, currentRip, lastImp->name);
+        sandbox->SetLastImpRead(0, nullptr);
+    } else {
+        for (auto module : sandbox->GetModuleList()) {
+            for (auto item : module->export_function) {
+                const auto vmAddress = module->base + item->function_address;
+                if (vmAddress == currentRip) {
+                    printf("[!!!]detect no correct call, currentRip: 0x%llx\n",
+                           currentRip);
+                    sandbox->SetLastImpRead(0, nullptr);
 
-    for (auto module : sandbox->GetModuleList()) {
+                    sandbox->EmulateApi(uc, vmAddress, currentRip, item->name);
-        for (auto item : module->export_function) {
+                }
-            const auto vmAddress = module->base + item->function_address;
-            if (vmAddress == currentRip) {
-                printf(
-                    "[%s] call function %s at file address: %llx lastRip: "
-                    "%llx\n",
-                    module->name, item->name, address, lastRip);
-                sandbox->EmulateApi(uc, vmAddress, currentRip, item->name);
             }
         }
     }
@@ -83,7 +93,7 @@ void handleCodeRun(uc_engine* uc, uint64_t address, uint32_t size,
         }
         cs_free(instruction, instructionCount);
 
-        //dumpVmenv(uc, userData);
+        // dumpVmenv(uc, userData);
     }
 }
 
@@ -126,7 +136,7 @@ void handleMemoryRead(uc_engine* uc, uc_mem_type type, uint64_t address,
         }
     }
 
-    uint64_t readAddress;
+    uint64_t readAddress = 0;
     auto readError =
         uc_mem_read(sandbox->GetUnicornHandle(), address, &readAddress, size);
     if (LOG_LEVEL > 0) {
@@ -135,6 +145,19 @@ void handleMemoryRead(uc_engine* uc, uc_mem_type type, uint64_t address,
             "ReadData: %p Rbp: %p\n",
             address, size, regRax, regRip, readError, readAddress, regRbp);
     }
+    for (auto imp : sandbox->GetImpFuncDict()) {
+        const auto vmAddress =
+            sandbox->GetPeInfo()->RecImageBase + imp->function_address;
+
+        if (vmAddress == address) {
+            printf(
+                "Handle ImpRead Address: [%s] call function %s at file "
+                "address: %llx readAddress: "
+                "%llx\n",
+                imp->name, imp->name, address, readAddress);
+            sandbox->SetLastImpRead(readAddress, imp);
+        }
+    }
 }
 void dumpVmenv(uc_engine* uc, void* userData) {
     auto* sandbox = static_cast<Sandbox*>(userData);
@@ -172,7 +195,7 @@ void dumpVmenv(uc_engine* uc, void* userData) {
                 sandbox->GetPeInfo()->isX64 ? UC_X86_REG_ECX : UC_X86_REG_ECX,
                 &Ecx);
     printf(
-        "[dumpVmenv] Rip: %p Rax: %p Rsp: %p Rbp: %p Rcx: %p Rdx: %p Eax: "
+        "[dumpVmenv] Rip: %p lastRip: %p Rax: %p Rsp: %p Rbp: %p Rcx: %p Rdx: %p Eax: "
         "%08x Ecx: %08x\n",
         Rip, Rax, Rsp, Rbp, Rcx, Rdx, Eax, Ecx);