From 8cfd24ab43c541779ba5fdc45f7702b103bcf5cb Mon Sep 17 00:00:00 2001 From: Huoji's <1296564236@qq.com> Date: Sun, 20 Apr 2025 23:43:54 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=B8=80=E4=B8=AA=E5=AF=BC?= =?UTF-8?q?=E8=87=B4=E5=B4=A9=E6=BA=83=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ai_anti_malware/ai_anti_malware.cpp | 46 +++++++++++++++++-- ai_anti_malware/ai_anti_malware.vcxproj | 3 +- .../ai_anti_malware.vcxproj.filters | 3 -- ai_anti_malware/head.h | 2 +- ai_anti_malware/ml.cpp | 19 +++++++- ai_anti_malware/sandbox.h | 1 + ml/malware_detector.cpp | 3 ++ 7 files changed, 66 insertions(+), 11 deletions(-) diff --git a/ai_anti_malware/ai_anti_malware.cpp b/ai_anti_malware/ai_anti_malware.cpp index 491fdbd..bb16ff2 100644 --- a/ai_anti_malware/ai_anti_malware.cpp +++ b/ai_anti_malware/ai_anti_malware.cpp @@ -23,10 +23,7 @@ auto getPeInfo(std::string inputFilePath) -> std::shared_ptr { 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 +void DetectMalwareInDirectory(const std::string& directoryPath) { + std::map 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; } diff --git a/ai_anti_malware/ai_anti_malware.vcxproj b/ai_anti_malware/ai_anti_malware.vcxproj index 013b624..aa47f0d 100644 --- a/ai_anti_malware/ai_anti_malware.vcxproj +++ b/ai_anti_malware/ai_anti_malware.vcxproj @@ -135,6 +135,8 @@ NDEBUG;_CONSOLE;_CRT_SECURE_NO_WARNINGS;%(PreprocessorDefinitions) true stdcpplatest + MaxSpeed + true Console @@ -191,7 +193,6 @@ - diff --git a/ai_anti_malware/ai_anti_malware.vcxproj.filters b/ai_anti_malware/ai_anti_malware.vcxproj.filters index 1592437..ceb4e86 100644 --- a/ai_anti_malware/ai_anti_malware.vcxproj.filters +++ b/ai_anti_malware/ai_anti_malware.vcxproj.filters @@ -191,9 +191,6 @@ 头文件\machine_learning - - 头文件\machine_learning - 头文件\sandbox diff --git a/ai_anti_malware/head.h b/ai_anti_malware/head.h index 09b0d16..6fb874e 100644 --- a/ai_anti_malware/head.h +++ b/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 diff --git a/ai_anti_malware/ml.cpp b/ai_anti_malware/ml.cpp index fb65f41..1b0766f 100644 --- a/ai_anti_malware/ml.cpp +++ b/ai_anti_malware/ml.cpp @@ -475,6 +475,24 @@ std::vector MachineLearning::ExtractFeatures(const uint8_t* buffer, // 提取所有特征 std::vector 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 propFeatures = @@ -512,7 +530,6 @@ std::vector MachineLearning::ExtractFeatures(const uint8_t* buffer, // 清理资源 peconv::free_pe_buffer(peBuffer); - return allFeatures; } diff --git a/ai_anti_malware/sandbox.h b/ai_anti_malware/sandbox.h index 277f9b7..ef1f959 100644 --- a/ai_anti_malware/sandbox.h +++ b/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 diff --git a/ml/malware_detector.cpp b/ml/malware_detector.cpp index 47a0971..a5cf307 100644 --- a/ml/malware_detector.cpp +++ b/ml/malware_detector.cpp @@ -1,5 +1,7 @@ #include #include +#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)