#include "util.h"

#include <iostream>
#include <locale>
#include <codecvt>

Uitil::Uitil()
{
}

std::wstring Uitil::stringToWString(const std::string &string)
{
    std::wstring_convert<std::codecvt_utf8<wchar_t>> cv;

    return cv.from_bytes(string);
}

HANDLE Uitil::getCurrentUserToken()
{
    // 查询sessionID
    PWTS_SESSION_INFO pSessionInfo = 0;
    DWORD dwCount = 0;
    ::WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE, 0, 1, &pSessionInfo, &dwCount);
    int session_id = 0;
    for (DWORD i = 0; i < dwCount; ++i)
    {
        WTS_SESSION_INFO si = pSessionInfo[i];
        if (WTSActive == si.State)
        {
            session_id = si.SessionId;
            break;
        }
    }

    ::WTSFreeMemory(pSessionInfo);

    // 查询token
    HANDLE current_token = 0;
    BOOL bRet = ::WTSQueryUserToken(session_id, &current_token);
    if (bRet == FALSE)
    {
        std::cout << "WTSQueryUserToken error, code:" << GetLastError() << std::endl;

        return nullptr;
    }

    HANDLE primaryToken = 0;
    bRet = ::DuplicateTokenEx(current_token, TOKEN_ASSIGN_PRIMARY | TOKEN_ALL_ACCESS, 0, SecurityImpersonation, TokenPrimary, &primaryToken);
    ::CloseHandle(current_token);
    if (bRet == FALSE)
    {
        std::cout << "DuplicateTokenEx error, code:" << GetLastError() << std::endl;

        return nullptr;
    }

    return primaryToken;
}

bool Uitil::runProgAsCurUser(HANDLE token, const std::string &progPath, const std::string &progArgs)
{
    STARTUPINFO StartupInfo = {0};
    PROCESS_INFORMATION processInfo;
    StartupInfo.cb = sizeof(STARTUPINFO);

    auto command = std::string("\"") + progPath + "\"";
    if (!progArgs.empty())
    {
        command += " " + progArgs;
    }

    void* lpEnvironment = NULL;
    BOOL resultEnv = ::CreateEnvironmentBlock(&lpEnvironment, token, FALSE);
    if (!resultEnv)
    {
        std::cout << "CreateEnvironmentBlock error, code:" << GetLastError() << std::endl;

        return false;
    }

    std::cout << "runProgAsCurUser, command:" << command << std::endl;

    // 获取到的hUnfilteredToken就是不受限的token，以token作为CreateProcessAsUser的第一个参数，就可以创建出具有管理员权限，并且属于当前用户的界面程序了，并且这种情况下，不需要加入窗口站。
    BOOL result = ::CreateProcessAsUser(token, 0, (LPSTR)(progPath.c_str()), NULL, NULL, FALSE, CREATE_NEW_CONSOLE | NORMAL_PRIORITY_CLASS | CREATE_UNICODE_ENVIRONMENT, lpEnvironment, 0, &StartupInfo, &processInfo);
    if(!result)
    {
        std::cout << "CreateProcessAsUser error, code:" << GetLastError() << std::endl;

        return false;
    }

    if(lpEnvironment != NULL)
    {
        ::DestroyEnvironmentBlock(lpEnvironment);
    }

    return true;
}

bool Uitil::RunProgAsCurUserAdminPrivilege(const std::string &progPath, const std::string &progArgs)
{
    // UAC开启时，当前用户拥有两个token，分别是受限的token和不受限的token。explorer.exe进程的token就属于受限的token。
    // 在服务程序中，可以用下面代码获取到受限的token。
    HANDLE primaryToken = getCurrentUserToken();
    if (primaryToken == 0)
    {
        std::cout << "GetCurrentUserToken error." << std::endl;

        return false;
    }

    // 由此token可以得到不受限的token
    bool isOpenOk = false;
    HANDLE hUnfilteredToken = NULL;
    DWORD dwSize = 0;
    BOOL bRet = ::GetTokenInformation(primaryToken, TokenLinkedToken, (VOID*)&hUnfilteredToken, sizeof(HANDLE), &dwSize);
    if (bRet)
    {
        isOpenOk = runProgAsCurUser(hUnfilteredToken, progPath, progArgs);

        ::CloseHandle(hUnfilteredToken);
    }
    else
    {
        // UAC未开时，继续使用原来的token打开
        std::cout << "GetTokenInformation error and continue open width primary token. code:" << GetLastError() << std::endl;

        isOpenOk = runProgAsCurUser(primaryToken, progPath, progArgs);
    }

    ::CloseHandle(primaryToken);

    return isOpenOk;
}