#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string>
#include <Userenv.h>
#include <Wtsapi32.h>
#pragma comment(lib, "WtsApi32.lib")
#pragma comment(lib, "advapi32.lib")
#pragma comment(lib, "userenv.lib")
using namespace std;

#define EPR                 fprintf(stderr,
#define ERR(str, chr)       if(opterr){EPR "%s%c\n", str, chr);}
int     opterr = 1;
int     optind = 1;
int    optopt;
char    *optarg;

int getopt (int argc, char *const argv[], char *opts)
{
    static int sp = 1;
    int c;
    char *cp;

    if (sp == 1)
        if (optind >= argc ||
           argv[optind][0] != '-' || argv[optind][1] == '\0')
            return -1;
        else if (strcmp(argv[optind], "--") == 0) {
            optind++;
            return -1;
        }
    optopt = c = argv[optind][sp];
    if (c == ':' || (cp=strchr(opts, c)) == 0) {
        ERR (": illegal option -- ", c);
        if (argv[optind][++sp] == '\0') {
            optind++;
            sp = 1;
        }
        return '?';
    }
    if (*++cp == ':') {
        if (argv[optind][sp+1] != '\0')
            optarg = &argv[optind++][sp+1];
        else if (++optind >= argc) {
            ERR (": option requires an argument -- ", c);
            sp = 1;
            return '?';
        } else
            optarg = argv[optind++];
        sp = 1;
    } else {
        if (argv[optind][++sp] == '\0') {
            sp = 1;
            optind++;
        }
        optarg = 0;
    }
    return c;
}

HANDLE GetUserToken(DWORD dwSessionId)
{
    HANDLE hImpersonationToken = 0;
    if (!WTSQueryUserToken(dwSessionId, &hImpersonationToken))
    {
        printf(" WTSQueryUserToken ERROR: %d\n", GetLastError());
        return FALSE;
    }
    DWORD dwNeededSize = 0;
    HANDLE *realToken = new HANDLE;
    TOKEN_USER *pTokenUser = NULL;
    PTOKEN_GROUPS pGroups = NULL;
    //twice call function
    if (!GetTokenInformation(hImpersonationToken, TokenUser, NULL, 0, &dwNeededSize))
    {
        if (GetLastError() == ERROR_INSUFFICIENT_BUFFER && dwNeededSize > 0)
        {
            pTokenUser = (TOKEN_USER*)new BYTE[dwNeededSize];
            if (!GetTokenInformation(hImpersonationToken, TokenUser, pTokenUser, dwNeededSize, &dwNeededSize))
            {
                printf("GetTokenInformation ERROR: %d", GetLastError());
            }
        }
        return hImpersonationToken;
    }
    return hImpersonationToken;
}

bool GetSessionUserName(DWORD dwSessionId, char username[256])
{
    LPTSTR pBuffer = NULL;
    DWORD dwBufferLen;
    if (!WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, dwSessionId, WTSUserName, &pBuffer, &dwBufferLen))
    {
        printf(" WTSQuerySessionInformation ERROR: %d\n", GetLastError());
        return FALSE;
    }
    lstrcpy(username ,pBuffer);
    WTSFreeMemory(pBuffer);
    return TRUE;
}

void Usage(void)
{
    printf("**********************usage******************\n");
    printf("USAGE: path:\\callsession.exe sxxxn path-exe-file arg1 arg2 arg3 ... \n");
    printf("**********************usage******************\n");
}

int main(int argc, char** argv)
{
    if(argc == 1)
    {
        Usage();
        return FALSE;
    }
    else if(argc > 1)
    {
        /* get opt comment if we need chengang882@pingan.com.cn at 2016-08-23
        int c;

        while ((c = getopt(argc, argv, "MNOVv+I:D:U:F:lg")) != -1)
            switch (c) {
            case 'N':

                break;
            case 'I':
                puts(optarg);

                break;
            case 'D':
            case 'U':
                puts(optarg);
                break;
            case 'M':

                break;
            case 'v':

                break;
            case 'V':

                break;
            case '+':

                break;
            default:
                break;
            }
        */
        int i;
        string lpCmdLine = "";
        LPSTR lpUsername;
        DWORD session_id = -1;
        DWORD session_count = 0;
        WTS_SESSION_INFOA *pSession = NULL;
        char username[256];
        BOOL blFound = FALSE;
        lpUsername = argv[1];
        // add all $2 $3 $4 $5 ... as second argument to program.
        for (i=2; i < argc; i++)
        {
            // printf("argument %d is %s.\n", i, argv[i]);
            lpCmdLine += argv[i];
            lpCmdLine += " ";
        }
        // printf("argument is %s\n", lpCmdLine);

        //EnumerateSessions
        if (!WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE, 0, 1, &pSession, &session_count))
        {
            printf("WTSEnumerateSessions ERROR: %d", GetLastError());
            return FALSE;
        }
        //Get the right user and his session id
        for(DWORD i = 0; i < session_count; ++i)
        {
            GetSessionUserName(pSession[i].SessionId,username);
            //if( (pSession[i].State == WTSActive) && (pSession[i].State != WTSDisconnected) )
            if(!strcmp(lpUsername, username))
            {
                printf("\tSession user name = %s\n",username);
                session_id = pSession[i].SessionId;
                printf("\tsession_id = %d\n",session_id);
                blFound = TRUE;
                break;
            }
        }
        if (!blFound){
            printf("No login username %s found.", lpUsername);
            return FALSE;
        }
        WTSFreeMemory(pSession); //free meme heap
        //Duplicate User Token
        HANDLE hTokenThis = GetUserToken(session_id);
        HANDLE hTokenDup = NULL;
        if (!DuplicateTokenEx(hTokenThis, TOKEN_ALL_ACCESS, NULL, SecurityIdentification, TokenPrimary, &hTokenDup))
        {
            printf("DuplicateTokenEx ERROR: %d\n", GetLastError());
            return FALSE;
        }
        if (!SetTokenInformation(hTokenDup, TokenSessionId, &session_id, sizeof(DWORD)))
        {
             printf("SetTokenInformation Error === %d\n",GetLastError());
             return FALSE;
        }
        //init this process info
        STARTUPINFO si;
        PROCESS_INFORMATION pi;
        ZeroMemory(&si, sizeof(STARTUPINFO));
        ZeroMemory(&pi, sizeof(PROCESS_INFORMATION));
        si.cb = sizeof(STARTUPINFO);
        si.lpDesktop = "WinSta0\\Default";
        //LPVOID pEnv = NULL;
        DWORD dwCreationFlag = NORMAL_PRIORITY_CLASS | CREATE_NEW_CONSOLE;
        //CreateEnvironmentBlock(&pEnv, hTokenDup, FALSE);
        if (!CreateProcessAsUser(hTokenDup, NULL, LPSTR(lpCmdLine.c_str()), NULL, NULL, FALSE, dwCreationFlag, NULL, NULL, &si, &pi))
        {
            printf("CreateProcessAsUser Error === %d\n",GetLastError());
            return FALSE;
        }
        printf("%s--%s--%s", lpUsername, LPSTR(lpCmdLine.c_str()), "OK");
    }

    return 0;
}