Feature: Client running as SYSTEM and support remote control

client/session.cpp

@@ -0,0 +1,176 @@
+#include "session.h"
+#include <wtsapi32.h>
+#include <userenv.h>
+#include <wtsapi32.h>
+#include <stdio.h>
+#include "common/logger.h"
+#pragma comment(lib, "Wtsapi32.lib")
+#pragma comment(lib, "Userenv.lib")
+#define SAFE_CLOSE_HANDLE(h) do{if((h)!=NULL&&(h)!=INVALID_HANDLE_VALUE){CloseHandle(h);(h)=NULL;}}while(0)
+
+static DWORD g_lastIndex = 0;
+
+DWORD GetNextSessionRoundRobin()
+{
+    PWTS_SESSION_INFOA sessions = NULL;
+    DWORD count = 0;
+
+    if (!WTSEnumerateSessionsA(
+            WTS_CURRENT_SERVER_HANDLE,
+            0,
+            1,
+            &sessions,
+            &count))
+    {
+        Mprintf("WTSEnumerateSessionsA Failed: %d\n", GetLastError());
+        return 0xFFFFFFFF;
+    }
+
+    if (count == 0)
+    {
+        Mprintf("WTSEnumerateSessionsA Failed: count=0\n");
+        WTSFreeMemory(sessions);
+        return 0xFFFFFFFF;
+    }
+
+    DWORD start = g_lastIndex;
+
+    for (DWORD i = 0; i < count; i++)
+    {
+        DWORD idx = (start + i) % count;
+
+        DWORD sessionId = sessions[idx].SessionId;
+
+        // 过滤无效 session
+        if (sessionId == 0)
+            continue;
+
+        WTS_CONNECTSTATE_CLASS* state = NULL;
+        DWORD bytes = 0;
+
+        if (WTSQuerySessionInformationA(
+                WTS_CURRENT_SERVER_HANDLE,
+                sessionId,
+                WTSConnectState,
+                (LPSTR*)&state,
+                &bytes))
+        {
+            BOOL ok = (state && (
+                *state == WTSActive ||
+                *state == WTSConnected));
+
+            WTSFreeMemory(state);
+
+            if (ok)
+            {
+                g_lastIndex = (idx + 1) % count;
+                DWORD result = sessionId;
+                WTSFreeMemory(sessions);
+                Mprintf("GetNextSessionRoundRobin Succeed: session=%d\n", result);
+                return result;
+            }
+        }
+    }
+
+    WTSFreeMemory(sessions);
+    return 0xFFFFFFFF;
+}
+
+BOOL StartProcessInSessionA(DWORD sessionId, const char* exePath, BYTE cmd)
+{
+    HANDLE hToken = NULL;
+    HANDLE hDupToken = NULL;
+    // 获取当前服务进程的 SYSTEM 令牌
+    char buf[500];
+    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_DUPLICATE | TOKEN_QUERY, &hToken)) {
+        sprintf(buf, "OpenProcessToken failed: %d\n", (int)GetLastError());
+        Mprintf(buf);
+        return FALSE;
+    }
+
+    // 复制为可用于创建进程的主令牌
+    if (!DuplicateTokenEx(hToken, MAXIMUM_ALLOWED, NULL,
+        SecurityImpersonation, TokenPrimary, &hDupToken)) {
+        sprintf(buf, "DuplicateTokenEx failed: %d\n", (int)GetLastError());
+        Mprintf(buf);
+        SAFE_CLOSE_HANDLE(hToken);
+        return FALSE;
+    }
+
+    // 修改令牌的会话 ID 为目标用户会话
+    if (!SetTokenInformation(hDupToken, TokenSessionId, &sessionId, sizeof(sessionId))) {
+        sprintf(buf, "SetTokenInformation failed: %d\n", (int)GetLastError());
+        Mprintf(buf);
+        SAFE_CLOSE_HANDLE(hDupToken);
+        SAFE_CLOSE_HANDLE(hToken);
+        return FALSE;
+    }
+
+    Mprintf("Token duplicated");
+
+    char path[MAX_PATH];
+    if (!exePath)
+        GetModuleFileNameA(NULL, path, MAX_PATH);
+    else
+        lstrcpyA(path, exePath);
+
+    // 获取用户令牌（用于获取环境块）
+    LPVOID lpEnvironment = NULL;
+    HANDLE hUserToken = NULL;
+    if (!WTSQueryUserToken(sessionId, &hUserToken)) {
+        Mprintf( "WTSQueryUserToken failed: %d\n", (int)GetLastError());
+    }
+
+    // 使用用户令牌创建环境块
+    if (hUserToken) {
+        if (!CreateEnvironmentBlock(&lpEnvironment, hUserToken, FALSE)) {
+            Mprintf("CreateEnvironmentBlock failed: %d\n", GetLastError());
+        }
+        CloseHandle(hUserToken);
+    }
+
+    STARTUPINFOA si = { 0 };
+    PROCESS_INFORMATION pi = { 0 };
+
+    si.cb = sizeof(si);
+    si.lpDesktop = (LPSTR)"winsta0\\default";
+    char cmdStr[300];
+    sprintf(cmdStr, "\"%s\" -cmd=%d", path, int(cmd));
+    BOOL result = CreateProcessAsUserA(
+        hDupToken,
+        NULL,
+        (LPSTR)cmdStr,
+        NULL,
+        NULL,
+        FALSE,
+        NORMAL_PRIORITY_CLASS | CREATE_NO_WINDOW | CREATE_UNICODE_ENVIRONMENT,
+        lpEnvironment,
+        NULL,
+        &si,
+        &pi);
+
+    if (result)
+    {
+        Mprintf("CreateProcessAsUserA Succeed\n");
+        CloseHandle(pi.hThread);
+        CloseHandle(pi.hProcess);
+    }
+    else {
+        Mprintf("CreateProcessAsUserA Failed [%d]: %s\n", GetLastError(), cmdStr);
+    }
+
+    CloseHandle(hDupToken);
+    CloseHandle(hToken);
+
+    return result;
+}
+
+void RunRoundRobinAgent(BYTE cmd)
+{
+    DWORD sessionId = GetNextSessionRoundRobin();
+
+    if (sessionId == 0xFFFFFFFF)
+        return;
+
+    StartProcessInSessionA(sessionId, NULL, cmd);
+}