update

ai_anti_malware/ai_anti_malware.cpp

@@ -313,24 +313,26 @@ auto doMalwareScan(int argc, char* argv[]) -> void {
 }
 
 int doSandbox(int argc, char* argv[]) {
-
-    std::string filePath = "Z:\\wangkun_muma\\opengl32.dll.bin";
-    auto peInfo = getPeInfo(filePath);
-    if (peInfo == nullptr) {
-        std::cout << "无法加载PE文件: " << filePath << std::endl;
-        return 1;
+    if (argc < 3) {
+        std::cout << "用法: " << argv[0] << " <文件路径> <地址>" << std::endl;
+        return;
     }
 
+    std::string filePath = argv[1];
+    auto peInfo = getPeInfo(filePath);
+    if (peInfo == nullptr) {
+        return 0;
+    }
     Sandbox se;
     se.InitEnv(peInfo);
-    se.Run(0x10002F20);
+    se.Run();
     return 0;
 }
 
 int main(int argc, char* argv[]) {
     // doMl(argc, argv);
     // doPredict(argc, argv);
-    // doMalwareScan(argc, argv);
-    doSandbox(argc, argv);
+    doMalwareScan(argc, argv);
+    // doSandbox(argc, argv);
     return 0;
 }

ai_anti_malware/sandbox.cpp

@@ -384,6 +384,7 @@ auto Sandbox::CreateModuleInfo(const char* dllName, uint64_t moduleBase,
 
 auto Sandbox::ResolveExport(uint64_t moduleBase)
     -> std::vector<std::shared_ptr<moudle_export>> {
+    std::vector<std::shared_ptr<moudle_export>> export_list;
     DWORD exportSize = 0;
     static RtlImageDirectoryEntryToDataFn fnRtlImageDirectoryEntryToData;
     if (fnRtlImageDirectoryEntryToData == nullptr) {
@@ -425,11 +426,11 @@ auto Sandbox::ResolveExport(uint64_t moduleBase)
             exportData.ordinal = static_cast<WORD>(
                 addressOfNameOrdinals[i] + exportDirectory->Base);  // 设置序号
 
-            m_exportFuncDict.push_back(
+            export_list.push_back(
                 std::make_shared<moudle_export>(exportData));
         }
     }
-    return m_exportFuncDict;
+    return export_list;
 }
 auto Sandbox::ResolveImportExports() -> void {
     for (auto module : m_moduleList) {
@@ -437,14 +438,14 @@ auto Sandbox::ResolveImportExports() -> void {
             continue;
         }
 
-        const auto exports = ResolveExport(module->real_base);
-        for (const auto item : exports) {
+        module->export_function = ResolveExport(module->real_base);
+        for (const auto item : module->export_function) {
             if (LOG_LEVEL > 0) {
                 printf("[ResolveImportExports] import export: [%s] %s => %llx\n", module->name,
                        item->name, item->function_address);
             }
 
-            module->export_function.push_back(item);
+            m_exportFuncDict.push_back(item);
         }
     }
 }