#include "precompile.h"
#include "toolkit.h"
#include <winpr/windows.h>
#include "memutil.h"

#ifdef _WIN32
#include <tlhelp32.h>    //CreateToolhelp32Snapshot
#endif //_WIN32

#define TAG TOOLKIT_TAG("process")

TOOLKIT_API int process_init(tk_process_t* proc)
{
	proc->handle = NULL;
	proc->pid = 0;
	return 0;
}

TOOLKIT_API int process_spawn(const tk_process_option_t* option, tk_process_t** proc)
{
	tk_process_t* new_process = NULL;
	int ret = 0;
	LPVOID pEnv = NULL;
	STARTUPINFOA si = { sizeof(STARTUPINFOA) };
	PROCESS_INFORMATION pi;
	HANDLE hProcess = NULL;
#if _WIN32
	DWORD dwSessionId;
	HANDLE hUserTokenDup, hThisToken;
	
	new_process = MALLOC_T(tk_process_t);
	if (new_process == NULL) {
		return -1;
	}
	dwSessionId = WTSGetActiveConsoleSessionId();
	if (OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, &hThisToken)) {
		LUID luid;
		TOKEN_PRIVILEGES tp;
		LookupPrivilegeValueA(NULL, SE_DEBUG_NAME, &luid);
		tp.PrivilegeCount = 1;
		tp.Privileges[0].Luid = luid;
		tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
		DuplicateTokenEx(hThisToken, MAXIMUM_ALLOWED, NULL,
			SecurityIdentification, TokenPrimary, &hUserTokenDup);
		SetTokenInformation(hUserTokenDup,
			TokenSessionId, (void*)dwSessionId, sizeof(DWORD));
		AdjustTokenPrivileges(hUserTokenDup, FALSE, &tp, sizeof(TOKEN_PRIVILEGES),
			(PTOKEN_PRIVILEGES)NULL, NULL);
		if (CreateProcessAsUserA(
			hUserTokenDup
			, option->file
			, option->params
			, NULL
			, NULL
			, FALSE
			, 0
			, pEnv
			, NULL
			, &si
			, &pi)) {
			CloseHandle(pi.hThread);
			new_process->pid = pi.dwProcessId;
			new_process->handle = pi.hProcess;
		}
		else {
			WLog_ERR(TAG, "create process as user failed.");
			FREE(new_process);
			ret = -1;
		}
		CloseHandle(hUserTokenDup);
		CloseHandle(hThisToken);
	}
	else {
		WLog_ERR(TAG, "open process token failed.");
		FREE(new_process);
		ret = -1;
	}
#else
	new_process = MALLOC_T(tk_process_t);
	if (new_process == NULL) {
		return -1;
	}
	if (CreateProcessA(
		option->file
		, option->params
		, NULL
		, NULL
		, FALSE
		, 0
		, pEnv
		, NULL
		, &si
		, &pi)) {
		CloseHandle(pi.hThread);
		new_process->pid = pi.dwProcessId;
		new_process->handle = pi.hProcess;
	}
	else {
		WLog_ERR(TAG, "create process  failed, GetLastError(%d).", GetLastError());
		FREE(new_process);
		ret = -1;
	}
#endif
	if (ret == 0) {
		*proc = new_process;
	}
	return ret;
}

TOOLKIT_API int process_compare(const tk_process_t* proc1, const tk_process_t* proc2)
{
	if(proc1 == NULL && proc1 == proc2) {
		return 0;
	}
	if(proc1 == NULL || proc2 == NULL) {
		return -1;
	}
	if(proc1->handle == proc2->handle) {
		return 0;
	}
	if (proc1->pid == proc2->pid) {
		return 0;
	}
	return -1;
}

TOOLKIT_API void process_close(tk_process_t* proc)
{
	DWORD exitCode = 0;
	if(GetExitCodeProcess(proc->handle, &exitCode) && GetLastError() != ERROR_INVALID_HANDLE) {
		CloseHandle(proc->handle);
		proc->handle = NULL;
	}
	proc->pid = 0;
}

TOOLKIT_API int process_exist_or_not(int pid)
{
	int ret = 0;

#ifdef _WIN32
	HANDLE hProcessSnap;
	PROCESSENTRY32 pe32;
	// Take a snapshot of all processes in the system.
	hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hProcessSnap == INVALID_HANDLE_VALUE) {
		return -1;
	}
	pe32.dwSize = sizeof(PROCESSENTRY32);
	if (!Process32First(hProcessSnap, &pe32)) {
		CloseHandle(hProcessSnap);          // clean the snapshot object
		return -1;
	}
	do {
		if (pid == pe32.th32ProcessID) {
			ret = 1;
			break;
		}
	} while (Process32Next(hProcessSnap, &pe32));
	CloseHandle(hProcessSnap);
#endif //_WIN32

	return ret;
}