#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);
}