update

ai_anti_malware/ai_anti_malware.cpp

@@ -29,28 +29,100 @@ auto getPeInfo(std::string inputFilePath) -> std::shared_ptr<BasicPeInfo> {
         sampleInfo->RecImageBase +
         (sampleInfo->isX64 ? sampleInfo->ntHead64->OptionalHeader.SizeOfImage
                            : sampleInfo->ntHead32->OptionalHeader.SizeOfImage);
+    printf("Debug - Memory mapping parameters:\n");
+    printf("RecImageBase: 0x%llx\n", sampleInfo->RecImageBase);
+    printf("peSize: 0x%llx\n", sampleInfo->peSize);
+    printf("Page aligned base: 0x%llx\n", sampleInfo->RecImageBase & ~0xFFF);
+    printf("Page aligned size: 0x%llx\n",
+           (sampleInfo->peSize + 0xFFF) & ~0xFFF);
+    sampleInfo->RecImageBase = sampleInfo->RecImageBase & ~0xFFF;
+    sampleInfo->peSize = (sampleInfo->peSize + 0xFFF) & ~0xFFF;
     return sampleInfo;
 }
-int main() {
-    auto sampleInfo = getPeInfo("z:\\Console_Test.exe");
-    // auto sampleInfo = getPeInfo("C:\\ConsoleApplication1.exe");
-    printf("input new file %s \n", sampleInfo->inputFilePath);
-    printf("is x64: %d\n", sampleInfo->isX64);
-    printf("is relocated: %d\n", sampleInfo->isRelocated);
-    printf("RecImageBase: %llx\n", sampleInfo->RecImageBase);
-    auto sandbox = std::make_shared<Sandbox>();
-    sandbox->InitEnv(sampleInfo);
-    sandbox->Run();
-    auto [peBuffer, peSize] = sandbox->DumpPE();
-
-    if (peBuffer) {
-        printf("peBuffer: %p\n", peBuffer.get());
-        printf("peSize: %d\n", peSize);
-        // peconv::dump_to_file("z:\\dumped_main.exe", peBuffer.get(), peSize);
-        MachineLearning ml;
-        ml.ExtractFeatures(peBuffer.get(), peSize, "z:\\features.txt");
+int doMl(int argc, char* argv[]) {
+    // 检查命令行参数
+    if (argc < 3) {
+        std::cout << "用法: " << argv[0] << " <样本目录路径> <输出CSV路径>"
+                  << std::endl;
+        std::cout << "或者: " << argv[0]
+                  << " -single <单个文件路径> <输出CSV路径>" << std::endl;
+        return 1;
     }
-    peBuffer.release();
+    MachineLearning ml;
+
+    if (std::string(argv[1]) == "-single") {
+        // 处理单个文件
+        if (argc < 4) {
+            std::cout << "处理单个文件时需要提供文件路径和输出CSV路径"
+                      << std::endl;
+            return 1;
+        }
+
+        std::string filePath = argv[2];
+        std::string csvPath = argv[3];
+
+        // 读取文件
+        std::vector<uint8_t> buffer = ml.ReadFileToBuffer(filePath);
+        if (buffer.empty()) {
+            std::cerr << "无法读取文件: " << filePath << std::endl;
+            return 1;
+        }
+
+        // 提取特征
+        std::vector<double> features =
+            ml.ExtractFeatures(buffer.data(), buffer.size());
+        if (features.empty()) {
+            std::cerr << "无法从文件提取特征: " << filePath << std::endl;
+            return 1;
+        }
+
+        // 导出到CSV
+        if (!ml.ExportToCSV(features, csvPath)) {
+            std::cerr << "无法导出到CSV文件: " << csvPath << std::endl;
+            return 1;
+        }
+
+        std::cout << "成功处理文件并导出特征到: " << csvPath << std::endl;
+    } else {
+        // 处理目录
+        std::string dirPath = argv[1];
+        std::string csvPath = argv[2];
+
+        std::cout << "开始处理目录: " << dirPath << std::endl;
+        std::cout << "特征将导出到: " << csvPath << std::endl;
+
+        if (!ml.ProcessDirectory(dirPath, csvPath)) {
+            std::cerr << "处理目录时发生错误" << std::endl;
+            return 1;
+        }
+    }
+    return 0;
+};
+int main(int argc, char* argv[]) {
+    doMl(argc, argv);
+    /*
+ auto sampleInfo = getPeInfo(
+     "E:\\对战平台\\CrowAntiCheat\\CrowAntiCheat\\client\\Console_"
+     "Test\\Release\\Console_Test.exe");
+ // auto sampleInfo = getPeInfo("C:\\ConsoleApplication1.exe");
+ printf("input new file %s \n", sampleInfo->inputFilePath);
+ printf("is x64: %d\n", sampleInfo->isX64);
+ printf("is relocated: %d\n", sampleInfo->isRelocated);
+ printf("RecImageBase: %llx\n", sampleInfo->RecImageBase);
+ auto sandbox = std::make_shared<Sandbox>();
+ sandbox->InitEnv(sampleInfo);
+ sandbox->Run();
+ auto [peBuffer, peSize] = sandbox->DumpPE();
+
+ if (peBuffer) {
+     printf("peBuffer: %p\n", peBuffer.get());
+     printf("peSize: %d\n", peSize);
+     // peconv::dump_to_file("z:\\dumped_main.exe", peBuffer.get(), peSize);
+     MachineLearning ml;
+     ml.ExtractFeatures(peBuffer.get(), peSize);
+ }
+ peBuffer.release();
+ */
     system("pause");
     return 0;
 }