修复一个导致崩溃的问题

ai_anti_malware/ai_anti_malware.cpp

@@ -23,10 +23,7 @@ auto getPeInfo(std::string inputFilePath) -> std::shared_ptr<BasicPeInfo> {
     sampleInfo->ntHead64 = peconv::get_nt_hdrs64((BYTE*)sampleInfo->peBuffer);
     sampleInfo->ntHead32 = peconv::get_nt_hdrs32((BYTE*)sampleInfo->peBuffer);
     sampleInfo->isX64 = peconv::is64bit((BYTE*)sampleInfo->peBuffer);
-    sampleInfo->RecImageBase =
-        sampleInfo->isX64
-            ? (DWORD64)sampleInfo->ntHead64->OptionalHeader.ImageBase
-            : (DWORD)sampleInfo->ntHead32->OptionalHeader.ImageBase;
+    sampleInfo->RecImageBase = MAIN_MODULE_BASE;
     sampleInfo->isRelocated =
         peconv::relocate_module((BYTE*)sampleInfo->peBuffer, sampleInfo->peSize,
                                 sampleInfo->RecImageBase);
@@ -335,11 +332,50 @@ int doSandbox(int argc, char* argv[]) {
     }
     return 0;
 }
+#include <filesystem>
 
+void DetectMalwareInDirectory(const std::string& directoryPath) {
+    std::map<DetectEngineType, int> detectionCount;
+
+    for (const auto& entry : std::filesystem::recursive_directory_iterator(directoryPath)) {
+        if (!entry.is_regular_file()) {
+            continue;
+        }
+
+        std::string filePath = entry.path().string();
+        std::cout << "Processing: " << filePath << std::endl;
+        DetectEngine scanner;
+        DetectEngineType result = scanner.DetectMalware(filePath);
+        detectionCount[result]++;
+    }
+
+    // 输出统计结果
+    std::cout << "\nDetection Summary:\n";
+    for (const auto& pair : detectionCount) {
+        std::string name;
+        switch (pair.first) {
+        case DetectEngineType::kNone: name = "None"; break;
+        case DetectEngineType::kPeStruct: name = "PE Struct"; break;
+        case DetectEngineType::kMachineLearning: name = "Machine Learning"; break;
+        case DetectEngineType::kSandbox: name = "Sandbox"; break;
+        }
+        std::cout << "  " << name << ": " << pair.second << "\n";
+    }
+}
 int main(int argc, char* argv[]) {
     // doMl(argc, argv);
     // doPredict(argc, argv);
     // doMalwareScan(argc, argv);
-    doSandbox(argc, argv);
+    // doSandbox(argc, argv);
+    /*
+    if (argc < 3) {
+        std::cout << "用法: " << argv[0] << " <文件夹路径>" << std::endl;
+        return 0;
+    }
+    
+    std::string filePath = argv[1];
+    */
+    std::string filePath = "Z:\\malware";
+    DetectMalwareInDirectory(filePath);
     return 0;
 }