修复一个导致崩溃的问题

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->RecImageBase = MAIN_MODULE_BASE;
-        sampleInfo->isX64
-            ? (DWORD64)sampleInfo->ntHead64->OptionalHeader.ImageBase
-            : (DWORD)sampleInfo->ntHead32->OptionalHeader.ImageBase;
     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;
 }

ai_anti_malware/ai_anti_malware.vcxproj

@@ -135,6 +135,8 @@
       <PreprocessorDefinitions>NDEBUG;_CONSOLE;_CRT_SECURE_NO_WARNINGS;%(PreprocessorDefinitions)</PreprocessorDefinitions>
       <ConformanceMode>true</ConformanceMode>
       <LanguageStandard>stdcpplatest</LanguageStandard>
+      <Optimization>MaxSpeed</Optimization>
+      <WholeProgramOptimization>true</WholeProgramOptimization>
     </ClCompile>
     <Link>
       <SubSystem>Console</SubSystem>
@@ -191,7 +193,6 @@
     <ClCompile Include="sandbox_malware_check.cpp" />
   </ItemGroup>
   <ItemGroup>
-    <ClInclude Include="..\ml\malware_detector.h" />
     <ClInclude Include="head.h" />
     <ClInclude Include="libpeconv\libpeconv\src\fix_dot_net_ep.h" />
     <ClInclude Include="libpeconv\libpeconv\src\ntddk.h" />

ai_anti_malware/ai_anti_malware.vcxproj.filters

@@ -191,9 +191,6 @@
     <ClInclude Include="ml.h">
       <Filter>头文件\machine_learning</Filter>
     </ClInclude>
-    <ClInclude Include="..\ml\malware_detector.h">
-      <Filter>头文件\machine_learning</Filter>
-    </ClInclude>
     <ClInclude Include="sandbox_api_winhttp.h">
       <Filter>头文件\sandbox</Filter>
     </ClInclude>

ai_anti_malware/head.h

@@ -1,5 +1,5 @@
 #pragma once
-#define LOG_LEVEL 1
+#define LOG_LEVEL 0
 
 #define _CRT_SECURE_NO_WARNINGS
 #include <iostream>

ai_anti_malware/ml.cpp

@@ -475,6 +475,24 @@ std::vector<double> MachineLearning::ExtractFeatures(const uint8_t* buffer,
 
     // 提取所有特征
     std::vector<double> allFeatures;
+    const size_t EXPECTED_PROPERTY_FEATURES = 14;  // 14个布尔值属性
+    const size_t EXPECTED_LIBRARY_FEATURES = 150;  // _libraries数组大小
+    const size_t EXPECTED_ENTROPY_FEATURES = 1;    // 文件熵
+    const size_t EXPECTED_ENTRYPOINT_FEATURES = 64; // EncodeEntrypoint实际使用64字节
+    const size_t EXPECTED_SECTION_FEATURES = 5;    // EncodeSections实际返回5个特征
+    const size_t EXPECTED_RATIO_FEATURES = 1;      // 代码比率
+    const size_t EXPECTED_SECTION_COUNT_FEATURES = 1; // 节区数量
+
+    const size_t TOTAL_EXPECTED_FEATURES =
+        EXPECTED_PROPERTY_FEATURES +
+        EXPECTED_LIBRARY_FEATURES +
+        EXPECTED_ENTROPY_FEATURES +
+        EXPECTED_ENTRYPOINT_FEATURES +
+        EXPECTED_SECTION_FEATURES +
+        EXPECTED_RATIO_FEATURES +
+        EXPECTED_SECTION_COUNT_FEATURES;
+
+    allFeatures.reserve(TOTAL_EXPECTED_FEATURES);
 
     // 1. PE段属性
     std::vector<double> propFeatures =
@@ -512,7 +530,6 @@ std::vector<double> MachineLearning::ExtractFeatures(const uint8_t* buffer,
 
     // 清理资源
     peconv::free_pe_buffer(peBuffer);
-
     return allFeatures;
 }
 

ai_anti_malware/sandbox.h

@@ -20,6 +20,7 @@
 #define HEAP_SIZE_32 0x5000000
 #define ENV_BLOCK_BASE 0x50000
 #define DLL_MODULE_BASE 0x130000
+#define MAIN_MODULE_BASE 0xff0000
 
 #define PEB_BASE 0x90000
 #define TEB_BASE 0x90000

ml/malware_detector.cpp

@@ -1,5 +1,7 @@
 #include <math.h>
 #include <string.h>
+#pragma optimize("", off)
+
 double sigmoid(double x) {
     if (x < 0.0) {
         double z = exp(x);
@@ -6621,3 +6623,4 @@ double score(double* input) {
          var99));
     return var100;
 }
+#pragma optimize("", on)