admin/awesome_anti_virus_engine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor import library parsing with SEH-safe function and README update

Browse Source 
- Extract SEH-based import library parsing into a separate function for better modularity
- Add new processImportWithSEH and processImportWithSEH_Internal functions to handle import parsing safely
- Update README.md with more detailed project description and initial engine classification
- Improve error handling and exception management during PE import directory traversal
- Standardize library name processing with lowercase transformation
This commit is contained in:
Huoji's
2025-03-09 04:42:28 +08:00
parent 95df007dbf
commit 10c56952c6
2 changed files with 77 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
ai_anti_malware/ml.cpp
View File

@@ -243,6 +243,63 @@ bool MachineLearning::ParseRichHeader(const uint8_t* peBuffer,
return true;
}
// 添加一个C风格的函数处理SEH部分
auto processImportWithSEH_Internal(const uint8_t* buffer, size_t bufferSize,
char** libNames, size_t* libCount,
size_t maxLibs) -> BOOL {
__try {
// 懒得JB处理了,累了.这里是不安全的
size_t impRva = 0;
size_t count = 0;
IMAGE_DATA_DIRECTORY* impDir =
peconv::get_directory_entry(buffer, IMAGE_DIRECTORY_ENTRY_IMPORT);
if (impDir) {
impRva = impDir->VirtualAddress;
IMAGE_IMPORT_DESCRIPTOR* impDesc =
reinterpret_cast<IMAGE_IMPORT_DESCRIPTOR*>(
RvaToPtr(impRva, (BYTE*)buffer));
while (impDesc && impDesc->Name != 0 && count < maxLibs) {
char* libName = reinterpret_cast<char*>(
RvaToPtr(impDesc->Name, (BYTE*)buffer));
if (libName) {
libNames[count] = libName;
count++;
}
impDesc++;
}
*libCount = count;
return TRUE;
}
return FALSE;
} __except (EXCEPTION_EXECUTE_HANDLER) {
printf("skip file: (access violation)\n");
return FALSE;
}
}
auto processImportWithSEH(const uint8_t* buffer, size_t bufferSize,
std::vector<std::string>& importedLibraries) -> void {
const size_t MAX_LIBS = 1000; // 设置一个合理的最大值
char* libNames[MAX_LIBS] = {0};
size_t libCount = 0;
// 调用处理SEH的内部函数
if (processImportWithSEH_Internal(buffer, bufferSize, libNames, &libCount,
MAX_LIBS)) {
// 将结果转换为C++对象
for (size_t i = 0; i < libCount; i++) {
if (libNames[i]) {
std::string libNameStr = libNames[i];
std::transform(libNameStr.begin(), libNameStr.end(),
libNameStr.begin(),
[](unsigned char c) { return std::tolower(c); });
importedLibraries.push_back(libNameStr);
}
}
}
}
std::vector<double> MachineLearning::ExtractFeatures(const uint8_t* buffer,
size_t bufferSize) {
// 使用libpeconv解析PE文件
@@ -381,33 +438,7 @@ std::vector<double> MachineLearning::ExtractFeatures(const uint8_t* buffer,
// 获取导入DLL列表
if (peInfo.hasImports) {
__try {
// 懒得JB处理了,累了.这里是不安全的
size_t impRva = 0;
IMAGE_DATA_DIRECTORY* impDir = peconv::get_directory_entry(
peBuffer, IMAGE_DIRECTORY_ENTRY_IMPORT);
if (impDir) {
impRva = impDir->VirtualAddress;
IMAGE_IMPORT_DESCRIPTOR* impDesc =
reinterpret_cast<IMAGE_IMPORT_DESCRIPTOR*>(
RvaToPtr(impRva, peBuffer));
while (impDesc && impDesc->Name != 0) {
char* libName = reinterpret_cast<char*>(
RvaToPtr(impDesc->Name, peBuffer));
if (libName) {
std::string libNameStr = libName;
std::transform(libNameStr.begin(), libNameStr.end(),
libNameStr.begin(), [](unsigned char c) {
return std::tolower(c);
});
importedLibraries.push_back(libNameStr);
}
impDesc++;
}
}
} __except (EXCEPTION_EXECUTE_HANDLER) {
printf("skip file: (access violation)\n");
}
processImportWithSEH(peBuffer, bufferSize, importedLibraries);
}
// 获取节区信息