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

更新文件路径，优化结构体定义，调整日志级别，增加模块处理函数，修复导入逻辑，提升代码可读性和维护性。

Browse Source
This commit is contained in:
Huoji's
2025-03-22 01:12:47 +08:00
parent 8e2e0c19ec
commit 2f08a0264a
6 changed files with 87 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ai_anti_malware/ai_anti_malware.cpp
View File

@@ -314,7 +314,7 @@ auto doMalwareScan(int argc, char* argv[]) -> void {
int doSandbox(int argc, char* argv[]) {
std::string filePath = "C:\\wangkun_muma\\opengl32.dll.bin";
std::string filePath = "Z:\\wangkun_muma\\opengl32.dll.bin";
auto peInfo = getPeInfo(filePath);
if (peInfo == nullptr) {
std::cout << "无法加载PE文件: " << filePath << std::endl;

6
ai_anti_malware/libpeconv/libpeconv/src/imports_loader.cpp
View File

@@ -1,7 +1,7 @@
#include "peconv/imports_loader.h"
#include <iostream>
#include <string>
using namespace peconv;
class FillImportThunks : public ImportThunksCallback
@@ -44,8 +44,8 @@ protected:
std::cout << "raw ordinal: " << std::hex << raw_ordinal << std::endl;
#endif
//hProc = funcResolver->resolve_func(lib_name, MAKEINTRESOURCEA(raw_ordinal));
return true;
hProc = funcResolver->resolve_func(lib_name, (LPSTR)std::to_string((uint32_t)raw_ordinal).c_str());
//return true;
}
else {
PIMAGE_IMPORT_BY_NAME by_name = (PIMAGE_IMPORT_BY_NAME)((ULONGLONG)modulePtr + desc->u1.AddressOfData);

66
ai_anti_malware/native_struct.h
View File

@@ -912,38 +912,37 @@ union SegmentSelector {
static_assert(sizeof(SegmentSelector) == 2, "Size check");
#include <poppack.h>
typedef struct _STARTUPINFOW32 {
DWORD cb;
DWORD lpReserved;
DWORD lpDesktop;
DWORD lpTitle;
DWORD dwX;
DWORD dwY;
DWORD dwXSize;
DWORD dwYSize;
DWORD dwXCountChars;
DWORD dwYCountChars;
DWORD dwFillAttribute;
DWORD dwFlags;
WORD wShowWindow;
WORD cbReserved2;
DWORD lpReserved2;
DWORD hStdInput;
DWORD hStdOutput;
DWORD hStdError;
} STARTUPINFOW32, * LPSTARTUPINFOW32;
DWORD cb;
DWORD lpReserved;
DWORD lpDesktop;
DWORD lpTitle;
DWORD dwX;
DWORD dwY;
DWORD dwXSize;
DWORD dwYSize;
DWORD dwXCountChars;
DWORD dwYCountChars;
DWORD dwFillAttribute;
DWORD dwFlags;
WORD wShowWindow;
WORD cbReserved2;
DWORD lpReserved2;
DWORD hStdInput;
DWORD hStdOutput;
DWORD hStdError;
} STARTUPINFOW32, *LPSTARTUPINFOW32;
static_assert(sizeof(STARTUPINFOW32) == 68, "Size check");
typedef struct tagPROCESSENTRY32W_32
{
DWORD dwSize;
DWORD cntUsage;
DWORD th32ProcessID; // this process
typedef struct tagPROCESSENTRY32W_32 {
DWORD dwSize;
DWORD cntUsage;
DWORD th32ProcessID; // this process
DWORD th32DefaultHeapID;
DWORD th32ModuleID; // associated exe
DWORD cntThreads;
DWORD th32ParentProcessID; // this process's parent process
LONG pcPriClassBase; // Base priority of process's threads
DWORD dwFlags;
WCHAR szExeFile[MAX_PATH]; // Path
DWORD th32ModuleID; // associated exe
DWORD cntThreads;
DWORD th32ParentProcessID; // this process's parent process
LONG pcPriClassBase; // Base priority of process's threads
DWORD dwFlags;
WCHAR szExeFile[MAX_PATH]; // Path
} PROCESSENTRY32W_32;
static_assert(sizeof(PROCESSENTRY32W_32) == 556, "Size check");
#pragma pack(push, 8)
@@ -952,10 +951,10 @@ typedef struct _RTL_CRITICAL_SECTION32 {
DWORD DebugInfo;
LONG LockCount;
LONG RecursionCount;
DWORD OwningThread; // from the thread's ClientId->UniqueThread
DWORD OwningThread; // from the thread's ClientId->UniqueThread
DWORD LockSemaphore;
DWORD SpinCount; // force size on 64-bit systems when packed
} RTL_CRITICAL_SECTION32, * PRTL_CRITICAL_SECTION32;
DWORD SpinCount; // force size on 64-bit systems when packed
} RTL_CRITICAL_SECTION32, *PRTL_CRITICAL_SECTION32;
#pragma pack(pop)
static_assert(sizeof(RTL_CRITICAL_SECTION32) == 24, "Size check");
@@ -1001,6 +1000,7 @@ struct moudle_import {
bool is_delayed_import;
};
struct moudle_import_ordinal {
std::string dll_name;
uint64_t function_address;
uint64_t ordinal;
};

62
ai_anti_malware/sandbox.cpp
View File

@@ -102,9 +102,10 @@ class ImportResolver : public peconv::t_function_resolver {
class cListImportNames : public peconv::ImportThunksCallback {
public:
cListImportNames(BYTE* _modulePtr, size_t _moduleSize,
std::vector<std::shared_ptr<moudle_import>>& name_to_addr,
std::vector<std::shared_ptr<moudle_import_ordinal>>& name_to_ordinal)
cListImportNames(
BYTE* _modulePtr, size_t _moduleSize,
std::vector<std::shared_ptr<moudle_import>>& name_to_addr,
std::vector<std::shared_ptr<moudle_import_ordinal>>& name_to_ordinal)
: ImportThunksCallback(_modulePtr, _moduleSize),
nameToAddr(name_to_addr),
ordinalImportFunc(name_to_ordinal) {}
@@ -149,11 +150,10 @@ class cListImportNames : public peconv::ImportThunksCallback {
import_data->function_address = call_via_rva;
import_data->is_delayed_import = false;
nameToAddr.push_back(import_data);
}
else {
} else {
auto importFunc = std::make_shared<moudle_import_ordinal>();
T_FIELD raw_ordinal = desc->u1.Ordinal & (~ordinal_flag);
importFunc->dll_name = lib_name;
importFunc->function_address = call_via_rva;
importFunc->ordinal = raw_ordinal;
ordinalImportFunc.push_back(importFunc);
@@ -181,8 +181,10 @@ class cFixImprot : public peconv::t_function_resolver {
if (strcmp(exp->name, func_name) == 0) {
auto newBase = reinterpret_cast<FARPROC>(
module->base + exp->function_address);
#ifdef LOG_LEVEL > 2
printf("fix import: %s => %llx \n", func_name, newBase);
// 返回在模拟器中的虚拟地址
#endif
return newBase;
}
}
@@ -192,15 +194,33 @@ class cFixImprot : public peconv::t_function_resolver {
// 如果没有找到精确匹配的模块名，尝试在所有模块中查找该函数
for (const auto& module : m_sandbox->m_moduleList) {
for (const auto& exp : module->export_function) {
auto newBase = reinterpret_cast<FARPROC>(
module->base + exp->function_address);
// 检查函数名是否匹配
if (strcmp(exp->name, func_name) == 0) {
auto newBase = reinterpret_cast<FARPROC>(
module->base + exp->function_address);
#ifdef LOG_LEVEL > 1
printf("fix import (fallback): %s found in %s => %llx \n",
func_name, module->name, newBase);
// 返回在模拟器中的虚拟地址
#endif
return newBase;
}
//序号导出,非常癌症的修复.
if (strcmp(module->name, lib_name) == 0) {
int ordinalNum = std::atoi(func_name);
if (exp->ordinal == ordinalNum) {
auto newBase = reinterpret_cast<FARPROC>(
module->base + exp->function_address);
#ifdef LOG_LEVEL > 1
printf("fix import (ordianal): %s found in [%s]%s => %llx \n",
func_name, module->name, exp->name, newBase);
// 返回在模拟器中的虚拟地址
#endif // LOG_LEVEL > 1
return newBase;
}
}
}
}
@@ -420,7 +440,7 @@ auto Sandbox::ResolveImportExports() -> void {
const auto exports = ResolveExport(module->real_base);
for (const auto item : exports) {
if (LOG_LEVEL > 0) {
printf("import export: [%s] %s => %llx\n", module->name,
printf("[ResolveImportExports] import export: [%s] %s => %llx\n", module->name,
item->name, item->function_address);
}
@@ -428,10 +448,9 @@ auto Sandbox::ResolveImportExports() -> void {
}
}
}
auto Sandbox::processImportModule(const moudle_import* importModule) -> void {
auto Sandbox::mapSystemModuleToVmByName(std::string systemName) -> void {
for (auto module : m_moduleList) {
if (strcmp(module->name, importModule->dll_name) == 0) {
if (strcmp(module->name, systemName.c_str()) == 0) {
if (LOG_LEVEL > 0) {
printf("skip module name: %s (already loaded)\n", module->name);
}
@@ -447,7 +466,7 @@ auto Sandbox::processImportModule(const moudle_import* importModule) -> void {
}
const std::string modulePath =
std::string(windowsPath) + systemDir + importModule->dll_name;
std::string(windowsPath) + systemDir + systemName;
// 加载PE模块
size_t mappedPeSize = 0;
@@ -459,14 +478,20 @@ auto Sandbox::processImportModule(const moudle_import* importModule) -> void {
}
// 添加到虚拟机
PushModuleToVM(importModule->dll_name, moduleBase);
PushModuleToVM(systemName.c_str(), moduleBase);
}
auto Sandbox::processImportModule(const moudle_import* importModule) -> void {
mapSystemModuleToVmByName(importModule->dll_name);
}
auto Sandbox::ResoveImport() -> void {
// 处理延迟导入
peconv::load_delayed_imports(static_cast<BYTE*>(m_peInfo->peBuffer), 0);
// 解析导入表
cListImportNames importCallback(static_cast<BYTE*>(m_peInfo->peBuffer), m_peInfo->peSize, m_impFuncDict, m_impFuncOrdinalDict);
cListImportNames importCallback(static_cast<BYTE*>(m_peInfo->peBuffer),
m_peInfo->peSize, m_impFuncDict,
m_impFuncOrdinalDict);
if (!peconv::process_import_table(static_cast<BYTE*>(m_peInfo->peBuffer),
m_peInfo->peSize, &importCallback)) {
@@ -477,6 +502,9 @@ auto Sandbox::ResoveImport() -> void {
for (const auto& importModule : m_impFuncDict) {
processImportModule(importModule.get());
}
for (const auto& importModule : m_impFuncOrdinalDict) {
mapSystemModuleToVmByName(importModule->dll_name);
}
}
auto Sandbox::SetupVirtualMachine() -> void {
SegmentSelector cs = {0};
@@ -717,6 +745,7 @@ auto Sandbox::InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void {
if (!peconv::load_imports(m_peInfo->peBuffer, &importFixer)) {
throw std::runtime_error("Failed to fix imports");
}
// 给所有导入表加c3
for (const auto& module : this->GetModuleList()) {
// 遍历导出函数查找对应名称
@@ -835,7 +864,7 @@ auto Sandbox::Run(uint64_t address) -> void {
std::cout << "Starting execution at " << std::hex << entryPoint
<< std::endl;
uint64_t timeout = 60 * 1000 * 1000;
uint64_t timeout = 2 * 60 * 1000 * 1000;
// 1.入口点是必须跑的
if (m_peInfo->isDll) {
// 给rcx和rdx设置dll应该设置的
@@ -1426,3 +1455,4 @@ auto Sandbox::AddModuleToLdr(const std::shared_ptr<struct_moudle>& module)
printf("Added module '%s' to LDR data tables at 0x%llx\n", module->name,
entryAddress);
}

1
ai_anti_malware/sandbox.h
View File

@@ -314,6 +314,7 @@ class Sandbox {
-> std::vector<std::shared_ptr<moudle_export>>;
auto InitApiHooks() -> void;
auto InitCommandLine(std::string commandLine) -> void;
auto mapSystemModuleToVmByName(std::string systemName) -> void;
std::vector<uint64_t> m_crossSectionExecution; // 记录跨区段执行地址
uint64_t m_lastExecuteSectionIndex = 0; // 上次执行的区段索引
uint64_t m_KSharedUserDataBase{0};

6
ai_anti_malware/sandbox_callbacks.cpp
View File

@@ -80,7 +80,7 @@ void handleCodeRun(uc_engine* uc, uint64_t address, uint32_t size,
}
}
lastRip = currentRip;
if (LOG_LEVEL > 0) {
if (LOG_LEVEL > 2) {
// 使用Capstone反汇编
cs_insn* instruction;
size_t instructionCount =
@@ -139,7 +139,7 @@ void handleMemoryRead(uc_engine* uc, uc_mem_type type, uint64_t address,
uint64_t readAddress = 0;
auto readError =
uc_mem_read(sandbox->GetUnicornHandle(), address, &readAddress, size);
if (LOG_LEVEL > 0) {
if (LOG_LEVEL > 2) {
printf(
"[handleMemoryRead] Address: %p Size: %p Rax: %p Rip: %p Error: %d "
"ReadData: %p Rbp: %p\n",
@@ -320,7 +320,7 @@ void handleMemoryWrite(uc_engine* uc, uc_mem_type type, uint64_t address,
}
}
if (LOG_LEVEL > 0) {
if (LOG_LEVEL > 2) {
printf("[handleMemoryWrite] Address: %p Size: %p Value: %p RIP: %p\n",
address, size, value, regRip);
}