update

ai_anti_malware/sandbox.cpp

@@ -1,5 +1,7 @@
 #include "sandbox.h"
+
 #include "sandbox_callbacks.h"
+
 std::string getDllNameFromApiSetMap(const std::string& apiSet) {
     const std::wstring wApiSet(apiSet.begin(), apiSet.end());
 
@@ -129,7 +131,37 @@ class cListImportNames : public peconv::ImportThunksCallback {
 
     std::vector<std::shared_ptr<moudle_import>>& nameToAddr;
 };
+class cFixImprot : public peconv::t_function_resolver {
+   public:
+    // 构造函数接收Sandbox实例的引用
+    explicit cFixImprot(Sandbox* sandbox) : m_sandbox(sandbox) {}
 
+    // 实现导入函数解析
+    virtual FARPROC resolve_func(LPSTR lib_name, LPSTR func_name) override {
+        // 遍历所有已加载的模块
+        for (const auto& module : m_sandbox->m_moduleList) {
+            // 检查模块名是否匹配
+            if (_stricmp(module->name, lib_name) == 0) {
+                // 遍历模块的导出函数
+                for (const auto& exp : m_sandbox->m_exportFuncDict) {
+                    // 检查函数名是否匹配
+                    if (strcmp(exp->name, func_name) == 0) {
+                        auto newBase = reinterpret_cast<FARPROC>(
+                            module->base + exp->function_address);
+                        printf("fix import: %s => %llx \n", func_name, newBase);
+                        // 返回在模拟器中的虚拟地址
+                        return newBase;
+                    }
+                }
+            }
+        }
+        __debugbreak();
+        return nullptr;
+    }
+
+   private:
+    Sandbox* m_sandbox;  // Sandbox实例的指针
+};
 Sandbox::Sandbox() {}
 
 Sandbox::~Sandbox() {}
@@ -167,7 +199,6 @@ auto Sandbox::PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void {
     newModule.real_base = moduleBase;
     newModule.entry = ntHeaders->OptionalHeader.AddressOfEntryPoint;
     newModule.size = ntHeaders->OptionalHeader.SizeOfImage;
-
     // 处理区段
     for (WORD i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++) {
         const auto& section = sectionHeader[i];
@@ -208,9 +239,13 @@ auto Sandbox::PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void {
     uc_mem_map(m_ucEngine, newModule.base, newModule.size,
                UC_PROT_READ | UC_PROT_EXEC);
     uc_mem_write(m_ucEngine, newModule.base, (void*)moduleBase, newModule.size);
+    if (peconv::relocate_module((BYTE*)moduleBase, newModule.size,
+                                newModule.base) == false) {
+        throw std::runtime_error("Failed to relocate module");
+    }
 }
-
-auto Sandbox::ResolveExport() -> void {
+auto Sandbox::ResolveExport(uint64_t moduleBase)
+    -> std::vector<std::shared_ptr<moudle_export>> {
     DWORD exportSize = 0;
     static RtlImageDirectoryEntryToDataFn fnRtlImageDirectoryEntryToData;
     if (fnRtlImageDirectoryEntryToData == nullptr) {
@@ -221,25 +256,25 @@ auto Sandbox::ResolveExport() -> void {
     // 获取导出表
     PIMAGE_EXPORT_DIRECTORY exportDirectory =
         static_cast<PIMAGE_EXPORT_DIRECTORY>(fnRtlImageDirectoryEntryToData(
-            m_peInfo->peBuffer, TRUE, IMAGE_DIRECTORY_ENTRY_EXPORT,
-            &exportSize));
+            reinterpret_cast<PUCHAR>(moduleBase), TRUE,
+            IMAGE_DIRECTORY_ENTRY_EXPORT, &exportSize));
 
     if (exportDirectory) {
         const DWORD numberOfNames = exportDirectory->NumberOfNames;
         PDWORD addressOfFunctions =
-            reinterpret_cast<PDWORD>(static_cast<PUCHAR>(m_peInfo->peBuffer) +
+            reinterpret_cast<PDWORD>(reinterpret_cast<PUCHAR>(moduleBase) +
                                      exportDirectory->AddressOfFunctions);
         PDWORD addressOfNames =
-            reinterpret_cast<PDWORD>(static_cast<PUCHAR>(m_peInfo->peBuffer) +
+            reinterpret_cast<PDWORD>(reinterpret_cast<PUCHAR>(moduleBase) +
                                      exportDirectory->AddressOfNames);
         PWORD addressOfNameOrdinals =
-            reinterpret_cast<PWORD>(static_cast<PUCHAR>(m_peInfo->peBuffer) +
+            reinterpret_cast<PWORD>(reinterpret_cast<PUCHAR>(moduleBase) +
                                     exportDirectory->AddressOfNameOrdinals);
 
         // 遍历导出函数
         for (size_t i = 0; i < numberOfNames; i++) {
             PCHAR functionName = reinterpret_cast<PCHAR>(
-                static_cast<PUCHAR>(m_peInfo->peBuffer) + addressOfNames[i]);
+                reinterpret_cast<PUCHAR>(moduleBase) + addressOfNames[i]);
 
             // 获取函数RVA
             const DWORD functionRva =
@@ -249,11 +284,24 @@ auto Sandbox::ResolveExport() -> void {
             moudle_export exportData{};
             memcpy(exportData.name, functionName, strlen(functionName));
             exportData.function_address = functionRva;
+            exportData.ordinal = static_cast<WORD>(
+                addressOfNameOrdinals[i] + exportDirectory->Base);  // 设置序号
 
             m_exportFuncDict.push_back(
                 std::make_shared<moudle_export>(exportData));
         }
     }
+    return m_exportFuncDict;
+}
+auto Sandbox::ResolveImportExports() -> void {
+    for (auto module : m_moduleList) {
+        const auto exports = ResolveExport(module->real_base);
+        for (const auto item : exports) {
+            printf("import export: [%s] %s => %llx\n", module->name, item->name,
+                   item->function_address);
+            module->export_function.push_back(item);
+        }
+    }
 }
 
 auto Sandbox::processImportModule(const moudle_import* importModule) -> void {
@@ -442,7 +490,6 @@ auto Sandbox::SetupVirtualMachine() -> void {
 }
 auto Sandbox::InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void {
     m_peInfo = peInfo;
-
     if (cs_open(CS_ARCH_X86, peInfo->isX64 ? CS_MODE_64 : CS_MODE_32,
                 &m_csHandle) != CS_ERR_OK) {
         throw std::runtime_error("Failed to initialize Capstone");
@@ -453,7 +500,14 @@ auto Sandbox::InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void {
         throw std::runtime_error("Failed to initialize Unicorn");
     }
     ResoveImport();
-    ResolveExport();
+    ResolveImportExports();
+
+    // 修复导入表
+    cFixImprot importFixer(this);
+    if (!peconv::load_imports(m_peInfo->peBuffer, &importFixer)) {
+        throw std::runtime_error("Failed to fix imports");
+    }
+
     uc_err ucErr = uc_mem_map(m_ucEngine, m_peInfo->RecImageBase,
                               m_peInfo->peSize, UC_PROT_ALL);
     if (ucErr != UC_ERR_OK) {
@@ -464,6 +518,8 @@ auto Sandbox::InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void {
     printf("map file to vm file: %llx\n", m_peInfo->RecImageBase);
     printf("map file to vm size: %llx\n", m_peInfo->peSize);
     SetupVirtualMachine();
+    InitCommandLine(peInfo->inputFilePath);
+
 }
 
 auto Sandbox::Run() -> void {
@@ -551,6 +607,7 @@ auto Sandbox::Run() -> void {
     // 开始执行
     std::cout << "Starting execution at " << std::hex << entryPoint
               << std::endl;
+    InitApiHooks();
     err = uc_emu_start(m_ucEngine, entryPoint, m_peInfo->imageEnd, 0, 0);
     if (err != UC_ERR_OK) {
         std::cerr << "Emulation error: " << uc_strerror(err) << std::endl;