update

ai_anti_malware/ai_anti_malware.cpp

@@ -311,16 +311,16 @@ auto doMalwareScan(int argc, char* argv[]) -> void {
     auto sampleType = scanner.DetectMalware(filePath);
     printf("sample type: %d \n", sampleType);
 }
-/*
+
 int doSandbox(int argc, char* argv[]) {
 
-    if (argc < 3) {
-        std::cout << "用法: " << argv[0] << " <文件路径> <地址>" << std::endl;
-        return;
-    }
-    std::string filePath = argv[1];
+    //if (argc < 3) {
+    //    std::cout << "用法: " << argv[0] << " <文件路径> <地址>" << std::endl;
+    //    return;
+    //}
+    //std::string filePath = argv[1];
 
-    std::string filePath = "C:\\mso.dll";
+    std::string filePath = "Z:\\mso.dll";
 
     auto peInfo = getPeInfo(filePath);
     if (peInfo == nullptr) {
@@ -331,11 +331,11 @@ int doSandbox(int argc, char* argv[]) {
     se.Run(0x180003980);
     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;
 }

ai_anti_malware/head.h

@@ -1,5 +1,5 @@
 #pragma once
-#define LOG_LEVEL 0
+#define LOG_LEVEL 1
 
 #define _CRT_SECURE_NO_WARNINGS
 #include <iostream>

ai_anti_malware/sandbox_api_emu.cpp

@@ -947,6 +947,8 @@ auto Sandbox::EmulateApi(uc_engine* uc, uint64_t address, uint64_t rip,
             rsp += 8;
             // 设置RIP为返回地址
             uc_reg_write(uc, UC_X86_REG_RIP, &return_address);
+            printf("ApiName set ip: %llx \n", return_address);
+
         } else {  // 32位系统
             // 读取4字节的返回地址
             uc_reg_read(uc, UC_X86_REG_ESP, &esp);

ai_anti_malware/sandbox_api_process.cpp

@@ -598,13 +598,13 @@ auto Api_CreateProcessA(void* sandbox, uc_engine* uc, uint64_t address)
     if (lpProcessInformation != 0) {
         if (context->GetPeInfo()->isX64) {
             struct PROCESS_INFORMATION64 {
-                HANDLE64 hProcess;
-                HANDLE64 hThread;
+                HANDLE hProcess;
+                HANDLE hThread;
                 DWORD dwProcessId;
                 DWORD dwThreadId;
             } pi;
-            pi.hProcess = (HANDLE64)hProcess;
-            pi.hThread = (HANDLE64)hThread;
+            pi.hProcess = (HANDLE)hProcess;
+            pi.hThread = (HANDLE)hThread;
             pi.dwProcessId = processId;
             pi.dwThreadId = threadId;
             uc_mem_write(uc, lpProcessInformation, &pi, sizeof(pi));

ai_anti_malware/sandbox_api_stl.cpp

@@ -1139,20 +1139,20 @@ auto Api_HeapAlloc(void* sandbox, uc_engine* uc, uint64_t address) -> void {
         dwBytes = temp_bytes;
     }
     // 这里如果想检查有效,得先跑main,再跑其他的,浪费时间了,操
-    
+
     // 检查堆句柄是否有效
-    /*
+
     uint64_t expected_handle =
         context->GetPeInfo()->isX64 ? HEAP_ADDRESS_64 : HEAP_ADDRESS_32;
     if (hHeap != expected_handle) {
         uint64_t null_ptr = 0;
         hHeap = expected_handle;
-        uc_reg_write(
-            uc, context->GetPeInfo()->isX64 ? UC_X86_REG_RAX : UC_X86_REG_EAX,
-            &null_ptr);
-        return;
+        // uc_reg_write(
+        //     uc, context->GetPeInfo()->isX64 ? UC_X86_REG_RAX :
+        //     UC_X86_REG_EAX, &null_ptr);
+        // return;
     }
-    */
+
     // 获取或创建堆段
     HeapSegment* segment = nullptr;
     auto it = context->m_heapSegments.find(hHeap);
@@ -2026,7 +2026,7 @@ auto Api_RtlFormatCurrentUserKeyPath(void* sandbox, uc_engine* uc,
 
     // 获取参数
     if (context->GetPeInfo()->isX64) {
-        // x64: rcx = KeyPathBuffer
+        // x64: rcx = KeyPathBuffer (PUNICODE_STRING)
         uc_reg_read(uc, UC_X86_REG_RCX, &keyPathBuffer);
     } else {
         // x86: 从栈上读取参数
@@ -2047,10 +2047,36 @@ auto Api_RtlFormatCurrentUserKeyPath(void* sandbox, uc_engine* uc,
     // 例如可以修改userKeyPath或在日志中记录查询操作
     printf("[*] RtlFormatCurrentUserKeyPath: 正在查询MCP服务器获取用户SID\n");
 
-    // 将路径写入到提供的缓冲区
-    if (keyPathBuffer != 0) {
-        uc_mem_write(uc, keyPathBuffer, userKeyPath,
-                     (wcslen(userKeyPath) + 1) * sizeof(wchar_t));
+    // 分配内存用于存储路径字符串
+    size_t pathLen = wcslen(userKeyPath);
+    size_t bufferSize = (pathLen + 1) * sizeof(wchar_t);
+    uint64_t stringBuffer = 0;
+
+    // 从堆中分配内存
+    if (context->GetPeInfo()->isX64) {
+        HeapSegment* segment = context->FindHeapSegment(HEAP_ADDRESS_64);
+        if (segment) {
+            stringBuffer = context->AllocateFromSegment(segment, bufferSize);
+        }
+    } else {
+        HeapSegment* segment = context->FindHeapSegment(HEAP_ADDRESS_32);
+        if (segment) {
+            stringBuffer = context->AllocateFromSegment(segment, bufferSize);
+        }
+    }
+
+    if (stringBuffer != 0 && keyPathBuffer != 0) {
+        // 将路径字符串写入到分配的缓冲区
+        uc_mem_write(uc, stringBuffer, userKeyPath, bufferSize);
+
+        // 创建UNICODE_STRING结构
+        UNICODE_STRING unicodeString;
+        unicodeString.Length = static_cast<USHORT>(pathLen * sizeof(wchar_t));
+        unicodeString.MaximumLength = static_cast<USHORT>(bufferSize);
+        unicodeString.Buffer = reinterpret_cast<PWSTR>(stringBuffer);
+
+        // 将UNICODE_STRING结构写入到提供的缓冲区
+        uc_mem_write(uc, keyPathBuffer, &unicodeString, sizeof(UNICODE_STRING));
     }
 
     // 返回NTSTATUS成功代码 (0x00000000 = STATUS_SUCCESS)
@@ -2059,8 +2085,10 @@ auto Api_RtlFormatCurrentUserKeyPath(void* sandbox, uc_engine* uc,
                  context->GetPeInfo()->isX64 ? UC_X86_REG_RAX : UC_X86_REG_EAX,
                  &status);
 
-    printf("[*] RtlFormatCurrentUserKeyPath: Buffer=0x%llx, Path=%ls\n",
-           keyPathBuffer, userKeyPath);
+    printf(
+        "[*] RtlFormatCurrentUserKeyPath: Buffer=0x%llx, StringBuffer=0x%llx, "
+        "Path=%ls\n",
+        keyPathBuffer, stringBuffer, userKeyPath);
 }
 
 // 添加FlsSetValue API实现