129 lines
4.7 KiB
C++
129 lines
4.7 KiB
C++
// ai_anti_malware.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
|
|
//
|
|
|
|
#include "head.h"
|
|
|
|
auto getPeInfo(std::string inputFilePath) -> std::shared_ptr<BasicPeInfo> {
|
|
auto sampleInfo = std::make_shared<BasicPeInfo>();
|
|
sampleInfo->inputFilePath = inputFilePath;
|
|
|
|
sampleInfo->peBuffer =
|
|
peconv::load_pe_module((const char*)sampleInfo->inputFilePath.c_str(),
|
|
sampleInfo->peSize, false, false);
|
|
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->isRelocated =
|
|
peconv::relocate_module((BYTE*)sampleInfo->peBuffer, sampleInfo->peSize,
|
|
sampleInfo->RecImageBase);
|
|
|
|
sampleInfo->entryPoint =
|
|
sampleInfo->isX64
|
|
? sampleInfo->ntHead64->OptionalHeader.AddressOfEntryPoint
|
|
: sampleInfo->ntHead32->OptionalHeader.AddressOfEntryPoint;
|
|
sampleInfo->imageEnd =
|
|
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 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;
|
|
}
|
|
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;
|
|
}
|