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

Add README and improve import library parsing with exception handling

Browse Source 
- Create initial README.md with project overview and motivation
- Add SEH-based exception handling in ExtractFeatures method for import library parsing
- Prevent potential access violations during PE import directory traversal
- Implement basic error logging for skipped files with access violations
This commit is contained in:
Huoji's
2025-03-09 04:24:46 +08:00
parent f80ba5d748
commit 95df007dbf
2 changed files with 38 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
ai_anti_malware/ml.cpp
View File

@@ -381,27 +381,32 @@ std::vector<double> MachineLearning::ExtractFeatures(const uint8_t* buffer,
// 获取导入DLL列表
if (peInfo.hasImports) {
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);
__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++;
}
impDesc++;
}
} __except (EXCEPTION_EXECUTE_HANDLER) {
printf("skip file: (access violation)\n");
}
}