diff --git a/ai_anti_malware/ai_anti_malware.cpp b/ai_anti_malware/ai_anti_malware.cpp index 49820bc..f47e2c3 100644 --- a/ai_anti_malware/ai_anti_malware.cpp +++ b/ai_anti_malware/ai_anti_malware.cpp @@ -313,24 +313,26 @@ auto doMalwareScan(int argc, char* argv[]) -> void { } int doSandbox(int argc, char* argv[]) { - - std::string filePath = "Z:\\wangkun_muma\\opengl32.dll.bin"; - auto peInfo = getPeInfo(filePath); - if (peInfo == nullptr) { - std::cout << "无法加载PE文件: " << filePath << std::endl; - return 1; + if (argc < 3) { + std::cout << "用法: " << argv[0] << " <文件路径> <地址>" << std::endl; + return; } + std::string filePath = argv[1]; + auto peInfo = getPeInfo(filePath); + if (peInfo == nullptr) { + return 0; + } Sandbox se; se.InitEnv(peInfo); - se.Run(0x10002F20); + se.Run(); return 0; } int main(int argc, char* argv[]) { // doMl(argc, argv); // doPredict(argc, argv); - // doMalwareScan(argc, argv); - doSandbox(argc, argv); + doMalwareScan(argc, argv); + // doSandbox(argc, argv); return 0; } diff --git a/ai_anti_malware/sandbox.cpp b/ai_anti_malware/sandbox.cpp index 3422d08..e6d4592 100644 --- a/ai_anti_malware/sandbox.cpp +++ b/ai_anti_malware/sandbox.cpp @@ -384,6 +384,7 @@ auto Sandbox::CreateModuleInfo(const char* dllName, uint64_t moduleBase, auto Sandbox::ResolveExport(uint64_t moduleBase) -> std::vector> { + std::vector> export_list; DWORD exportSize = 0; static RtlImageDirectoryEntryToDataFn fnRtlImageDirectoryEntryToData; if (fnRtlImageDirectoryEntryToData == nullptr) { @@ -425,11 +426,11 @@ auto Sandbox::ResolveExport(uint64_t moduleBase) exportData.ordinal = static_cast( addressOfNameOrdinals[i] + exportDirectory->Base); // 设置序号 - m_exportFuncDict.push_back( + export_list.push_back( std::make_shared(exportData)); } } - return m_exportFuncDict; + return export_list; } auto Sandbox::ResolveImportExports() -> void { for (auto module : m_moduleList) { @@ -437,14 +438,14 @@ auto Sandbox::ResolveImportExports() -> void { continue; } - const auto exports = ResolveExport(module->real_base); - for (const auto item : exports) { + module->export_function = ResolveExport(module->real_base); + for (const auto item : module->export_function) { if (LOG_LEVEL > 0) { printf("[ResolveImportExports] import export: [%s] %s => %llx\n", module->name, item->name, item->function_address); } - module->export_function.push_back(item); + m_exportFuncDict.push_back(item); } } }