#include "StdAfx.h"
#include "AutoFunc.h"
#include <Psapi.h>
#include <shellapi.h>

#pragma comment(lib, "shell32.lib") 

#pragma comment(lib, "Psapi.lib")


SystemCmds::SystemCmds(void)
{
	::CoInitialize(NULL); 

}


SystemCmds::~SystemCmds(void)
{
	CoUninitialize();
}

SystemCmds g_SystemCmds_Init;


BOOL SystemCmds::CreateTheProcess_Console(const TCHAR *pfullFilePath, PROCESS_INFORMATION &pinfo, bool HideWindow)
{
	string fullFilePath = pfullFilePath;
	size_t index = fullFilePath.rfind('\\');
	std::string strPathOnly = fullFilePath.substr(0, index);
	char szCommand[MAX_PATH];
	memset(szCommand, 0, MAX_PATH);
	strcpy_s(szCommand, MAX_PATH, pfullFilePath);

	STARTUPINFO si;
	ZeroMemory(&si, sizeof(si));
	si.cb = sizeof(si);
	ZeroMemory(&pinfo, sizeof(pinfo));

	if (HideWindow)
	{
		si.dwFlags = STARTF_USESHOWWINDOW;
		si.wShowWindow = SW_HIDE;
	}

	if (!CreateProcess(NULL,   // No module name (use command line)
		szCommand,        // Command line
		NULL,           // Process handle not inheritable
		NULL,           // Thread handle not inheritable
		FALSE,          // Set handle inheritance to FALSE
		CREATE_NEW_CONSOLE,              // No creation flags
		NULL,           // Use parent's environment block
		//NULL,
		strPathOnly.c_str(),           // Use parent's starting directory 
		&si,            // Pointer to STARTUPINFO structure
		&pinfo)           // Pointer to PROCESS_INFORMATION structure
		)
	{
		printf("CreateProcess failed (%d).\n", GetLastError());
		return FALSE;
	}


	return TRUE;

}


BOOL SystemCmds::RunTheProcess(const TCHAR *fullFilePath,bool HideWindow)
{
	INT ShowFlag = SW_SHOWNORMAL;
	if (HideWindow)
	{
		ShowFlag = SW_HIDE;
	}
	string fullFilePathL = fullFilePath;
	size_t index = fullFilePathL.rfind('\\');
	std::string strPathOnly = fullFilePathL.substr(0, index);
	std::string fileName = fullFilePathL.substr(index + 1,fullFilePathL.size() - (index + 1));
	if ((INT_PTR)ShellExecute(NULL, NULL, fileName.c_str(), NULL, strPathOnly.c_str(), ShowFlag) > 32)
	{
		return TRUE;
	}
	return FALSE;
}


DWORD SystemCmds::FindProcess(const TCHAR *strProcessName)
{
	DWORD aProcesses[4096], cbNeeded;
	HANDLE hProcess;
	TCHAR szProcessName[MAX_PATH];

	if ( !EnumProcesses( aProcesses, sizeof(aProcesses), &cbNeeded ) )  
	{
		return 0;
	}
	for(int i=0; i< (int) (cbNeeded / sizeof(DWORD)); i++)
	{
		hProcess = OpenProcess(  PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, aProcesses[i]);
		if(hProcess)
		{
			memset(szProcessName,0,sizeof(szProcessName));
			GetModuleFileNameEx(hProcess,NULL,szProcessName,MAX_PATH);
			CloseHandle(hProcess);
			string prc = szProcessName;
			if(prc.size() == 0)
			{
				continue;
			}
			std::size_t found = prc.find_last_of("\\");
			string filename = prc.substr(found+1);
			//_tprintf(_T("filename:%s\n"),filename.c_str());
			if (_stricmp(filename.c_str(), strProcessName) == 0)
			{
				return(aProcesses[i]);
			}
		}
	}

	return 0;
}

DWORD SystemCmds::KillProcess(const TCHAR *strProcessName)
{
	HANDLE TargetProcess = OpenProcess(PROCESS_TERMINATE, FALSE, SystemCmds::FindProcess(strProcessName));

	if(TargetProcess == NULL)
	{
		return 0;
	}

	TerminateProcess(TargetProcess, 0);

	CloseHandle(TargetProcess);

	return 1;
}


BOOL SystemCmds::EnableComputerCtrl()
{
	HANDLE tokenHandle;
	LUID privilegeLUID;
	PTOKEN_PRIVILEGES tokenPrivileges = new TOKEN_PRIVILEGES();
	PTOKEN_PRIVILEGES olePrivileges = new TOKEN_PRIVILEGES();
	try
	{
		if(!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &tokenHandle))
		{
			return FALSE;
		}
		if (!LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &privilegeLUID))
		{
			return FALSE;
		}
		tokenPrivileges->PrivilegeCount = 1;
		tokenPrivileges->Privileges->Attributes = SE_PRIVILEGE_ENABLED;
		tokenPrivileges->Privileges->Luid = privilegeLUID;
		DWORD dsize;
		if(!AdjustTokenPrivileges(tokenHandle, 0, tokenPrivileges, 4 + (12 * tokenPrivileges->PrivilegeCount), olePrivileges, &dsize))
		{
			return FALSE;
		}
	}
	catch(...)
	{
		return FALSE;
	}
	return TRUE;
}

BOOL SystemCmds::ShutDownComputer(BOOL Shutdown)
{
	BOOL ret = SystemCmds::EnableComputerCtrl();


	if (Shutdown == FALSE)
	{
		ret &= ExitWindowsEx(EWX_REBOOT | EWX_FORCE, 0);
	}
	else
	{
		ret &= ExitWindowsEx(EWX_POWEROFF | EWX_FORCE, 0);
	}

	return ret;

}