From 3a6e331f31538b2c5585c54512e8d0767a5a1408 Mon Sep 17 00:00:00 2001 From: Huoji's <1296564236@qq.com> Date: Thu, 20 Mar 2025 02:18:00 +0800 Subject: [PATCH] update --- .vscode/settings.json | 3 +- ai_anti_malware/ai_anti_malware.cpp | 2 +- ai_anti_malware/sandbox.cpp | 17 ++++++---- ai_anti_malware/sandbox.h | 17 ++++++++-- ai_anti_malware/sandbox_api_emu.cpp | 8 ++--- ai_anti_malware/sandbox_api_stl.cpp | 4 --- ai_anti_malware/sandbox_callbacks.cpp | 47 ++++++++++++++++++++------- 7 files changed, 66 insertions(+), 32 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index a17f25b..6349389 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -71,6 +71,7 @@ "stop_token": "cpp", "thread": "cpp", "xlocbuf": "cpp", - "xlocmes": "cpp" + "xlocmes": "cpp", + "unordered_set": "cpp" } } \ No newline at end of file diff --git a/ai_anti_malware/ai_anti_malware.cpp b/ai_anti_malware/ai_anti_malware.cpp index f67344c..4d84582 100644 --- a/ai_anti_malware/ai_anti_malware.cpp +++ b/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; diff --git a/ai_anti_malware/sandbox.cpp b/ai_anti_malware/sandbox.cpp index cb53f7d..aa33398 100644 --- a/ai_anti_malware/sandbox.cpp +++ b/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 { // 解析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(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 peInfo) -> void { } // 一定要确保他是第一个. auto newModule = - CreateModuleInfo("huoji.exe", m_peInfo->RecImageBase, + CreateModuleInfo("huoji.exe", m_peInfo->RecImageBase, m_peInfo->RecImageBase, reinterpret_cast(m_peInfo->peBuffer)); _ASSERTE(m_moduleList.size() == 0); m_moduleList.push_back(newModule); diff --git a/ai_anti_malware/sandbox.h b/ai_anti_malware/sandbox.h index 8b4b5b3..4e74666 100644 --- a/ai_anti_malware/sandbox.h +++ b/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 process_enum_state; // 用于跟踪每个句柄的枚举状态 - // Public methods auto InitEnv(std::shared_ptr 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> { + return m_impFuncDict; + } + auto GetLastImpRead() -> std::pair> { + return m_lastImpRead; + } + auto SetLastImpRead(uint64_t address, std::shared_ptr imp) -> void { + m_lastImpRead = { address, imp }; + } private: std::shared_ptr m_peInfo; + std::pair> 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; auto ResolveExport(uint64_t moduleBase) diff --git a/ai_anti_malware/sandbox_api_emu.cpp b/ai_anti_malware/sandbox_api_emu.cpp index 5e904f6..0417174 100644 --- a/ai_anti_malware/sandbox_api_emu.cpp +++ b/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); - return; + //uc_emu_stop(uc); + return false; } \ No newline at end of file diff --git a/ai_anti_malware/sandbox_api_stl.cpp b/ai_anti_malware/sandbox_api_stl.cpp index a8cb73b..03fc224 100644 --- a/ai_anti_malware/sandbox_api_stl.cpp +++ b/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); diff --git a/ai_anti_malware/sandbox_callbacks.cpp b/ai_anti_malware/sandbox_callbacks.cpp index 5b70e0b..8d68abd 100644 --- a/ai_anti_malware/sandbox_callbacks.cpp +++ b/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()) { - 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); + 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(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);