#include "StdAfx.h"
#include "PythonExecutor.h"
#include <fstream>
#include <io.h>
#include <fcntl.h>
#include "Python.h"
// 服务进程头文件;
#include <Wtsapi32.h>
#include <TlHelp32.h>
#pragma comment(lib, "wtsapi32.lib")
#include <Userenv.h>
#pragma comment(lib,"userenv.lib")


CPythonExecutor::CPythonExecutor(void)
{
	m_hWorkThread = NULL;
	m_hStdoutLogThread = NULL;
	m_hStderrLogThread = NULL;

	m_hStdOutRead = NULL;
	m_hStdOutWrite = NULL;
	m_hStdErrorRead = NULL;
	m_hStdErrorWrite = NULL;

	memset(m_szScriptPath, 0, MAX_PATH);
	memset(m_szLogPath, 0, MAX_PATH);
	memset(m_szExtraSentence, 0, MAX_PATH);

	memset(&m_si, 0, sizeof(m_si));
	memset(&m_pi, 0, sizeof(m_pi));
	m_si.cb = sizeof(m_si);

	m_bStopLogExport = FALSE;

	m_pCaseObj = NULL;
	m_ulStartTickCount = 0;
}

CPythonExecutor::~CPythonExecutor(void)
{
	TRACE1("\t~CPythonExecutor:%p\n\n", this);
	EndLogThread();
	EndWorkThread();	
}

DWORD CPythonExecutor::_WorkerThread(LPVOID lpParam)
{
	CPythonExecutor* that = (CPythonExecutor*)lpParam;

#if (defined _DEBUG && !defined DBG_SERVER) || defined RTEST
	that->RunScriptProcess();
#else
	that->ServiceRunScriptProcess();
#endif

	// 延时;
	Sleep(2000);
	// 结束日志线程;
	that->m_bStopLogExport = TRUE;
	// 可能出错：m_pCaseObj可能会提前释放了,所以需要保证不能提前释放;
	GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, _T("脚本执行线程结束,用例名=%s"), that->m_pCaseObj->strCaseName.c_str());

	return 0;
}

DWORD CPythonExecutor::_StdoutLogExportThread(LPVOID lpParam)
{
	CPythonExecutor* that = (CPythonExecutor*)lpParam;

	DWORD dwBytesToRead;
	CHAR chBuf[BUFSIZE] = {0};
	BOOL bSuccess = FALSE;

	do
	{
		// 预览管道中数据,有才ReadFile;
		PeekNamedPipe(that->m_hStdOutRead,chBuf,BUFSIZE,&dwBytesToRead,0,0);
		if( dwBytesToRead ) {
			bSuccess = ReadFile(that->m_hStdOutRead, chBuf, BUFSIZE, &dwBytesToRead, NULL);
			if (!bSuccess || dwBytesToRead == 0 || !_tcslen(chBuf)) 
				continue;

			that->m_ulStartTickCount = GetTickCount64();
			GLOBAL::WritePythonLog(that->m_szLogPath, chBuf);
			// 更新日志时间;
			if ( that->m_pCaseObj )
				that->m_pCaseObj->_ulStartTickCount = GetTickCount64();
			memset(chBuf, 0, BUFSIZE);
		}
		Sleep(5);
	} while (!that->m_bStopLogExport);

	GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, _T("脚本标准输出日志线程结束,用例名=%s"), that->m_pCaseObj->strCaseName.c_str());

	return 0;
}

DWORD CPythonExecutor::_StderrLogExportThread(LPVOID lpParam)
{
	CPythonExecutor* that = (CPythonExecutor*)lpParam;

	DWORD dwRead;
	CHAR chBuf[BUFSIZE] = {0};
	CHAR szMsg[BUFSIZE+MAX_PATH] = {0};
	BOOL bSuccess = FALSE;

	do
	{
		// 预览管道中数据,有才ReadFile;
		PeekNamedPipe(that->m_hStdErrorRead,chBuf,BUFSIZE,&dwRead,0,0);
		if( dwRead ) {
			bSuccess = ReadFile(that->m_hStdErrorRead, chBuf, BUFSIZE, &dwRead, NULL);
			if (!bSuccess || dwRead == 0 || !_tcslen(chBuf)) 
				continue;

			GLOBAL::WritePythonLog(that->m_szLogPath, chBuf);
			// 脚本结果异常;
			that->m_pCaseObj->_nExecutionResult = SATHTTP::ABNORMAL;

			_stprintf_s(szMsg, "_StderrLogExportThread: 脚本出现异常, 内容=%s\n", chBuf);
			OutputDebugString(szMsg);
			memset(chBuf, 0, BUFSIZE);
		}
		Sleep(5);
	} while (!that->m_bStopLogExport); 
	
	// 结束重定向;
	that->EndSubprocessStdOut();
	GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, _T("脚本标准错误日志线程结束,用例名=%s"), that->m_pCaseObj->strCaseName.c_str());

	return 0;
}

int CPythonExecutor::RedirectSubprocessStdout(LPSTARTUPINFO si /*=NULL*/)
{
	// Python脚本中,必须使用sys.__stdout__()
	SECURITY_ATTRIBUTES sa;
	sa.bInheritHandle = TRUE;
	sa.lpSecurityDescriptor = NULL;
	sa.nLength = sizeof(sa);

	// 创建stdout的管道;
	if (!CreatePipe(&m_hStdOutRead, &m_hStdOutWrite, &sa, 0)) {
		OutputDebugString("Error:创建stdout管道失败\n");
		return -1;
	}

#if 0// 标准输出和标准错误都由同一个句柄重定向,此时只需要一个线程来读取管道内容;
	// 创建stderr的管道,由于stderr一般就是stdout,直接复制句柄;
	if (!DuplicateHandle(GetCurrentProcess(), m_hStdOutWrite, GetCurrentProcess(), &m_hStdErrorWrite, 0, TRUE, DUPLICATE_SAME_ACCESS)) {
		OutputDebugString("创建stderr管道失败\n");
		return -2;
	}
#else
	// 标准输出和标准错误分开,需要二个线程来读取管道内容;
	if (!CreatePipe(&m_hStdErrorRead, &m_hStdErrorWrite, &sa, 0)) {
		OutputDebugString("Error:创建stderr管道失败\n");
		return -1;
	}
#endif

	si->dwFlags |= STARTF_USESTDHANDLES;
	// 将子进程的stdout输出到句柄hStdOutWrite;
	si->hStdOutput = m_hStdOutWrite;
	// 将子进程的stderr输出到句柄hStdErrWrite;
	si->hStdError = m_hStdErrorWrite;

	return 0;
}

int CPythonExecutor::RunScriptProcess()
{
	if (!PathFileExists(m_szScriptPath)) {
		GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "【Error】脚本路径无效:%s", m_szScriptPath);
		return -1;
	}

	// 初始化参数;
	::memset(&m_si, 0, sizeof(m_si));
	::memset(&m_pi, 0, sizeof(m_pi));
	m_si.cb = sizeof(m_si);
	GetStartupInfo(&m_si); 

	// 强制stdion, stdout和stderr完全无缓冲:python -u
	TCHAR szCommandLine[MAX_PATH] = { 0 };
	if (_tcslen(m_szExtraSentence))
		_stprintf_s(szCommandLine, _T("python -W ignore -u \"%s\" \"%s\""), m_szScriptPath, m_szExtraSentence);
	else
		_stprintf_s(szCommandLine, _T("python -W ignore -u \"%s\""), m_szScriptPath);

	// 重定向输出;
	RedirectSubprocessStdout(&m_si);
	// 恢复日志线程;
	ResumeThread(m_hStdoutLogThread);
	ResumeThread(m_hStderrLogThread);

	// 启动子进程;
	if (!CreateProcess(
		NULL,					// No module name (use command line)
		szCommandLine,			// Command line
		NULL,					// Process handle not inheritable
		NULL,					// Thread handle not inheritable
		TRUE,					// Set handle inheritance to TRUE
		0,						// No creation flags
		NULL,					// Use parent's environment block
		NULL,					// Use parent's starting directory 
		&m_si,					// Pointer to STARTUPINFO structure
		&m_pi)					// Pointer to PROCESS_INFORMATION structure
		)
	{
		GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "【Error】创建子进程失败 (%d)，", GetLastError());
		return -3;
	}

	// 记录进程ID;
	if ( m_pCaseObj )
		m_pCaseObj->__dwPythonPID = m_pi.dwProcessId;

	GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "Python进程ID=%ld", m_pi.dwProcessId);
	// 等待进程完成退出.
	WaitForSingleObject(m_pi.hProcess, INFINITE);

	// 必须等待，等待管道输出结果;
	Sleep(6000);
	GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "脚本进程结束（%ld）",  m_pi.dwProcessId);

	// 关闭管道前,先取消管道的IO操作;
	CancelIo(m_hStdOutWrite);
	CancelIo(m_hStdErrorWrite);
	// 同时需要关闭输出的管道,否则ReadFile会阻塞;
	CloseHandle(m_hStdOutWrite);
	m_hStdOutWrite = NULL;
	CloseHandle(m_hStdErrorWrite);
	m_hStdErrorWrite = NULL;

	// 关闭进程句柄. 
	CloseHandle(m_pi.hProcess);
	CloseHandle(m_pi.hThread);

	// 重置;
	memset(&m_si, 0, sizeof(m_si));
	memset(&m_pi, 0, sizeof(m_pi));
	m_si.cb = sizeof(m_si);

	return 0;
}

int CPythonExecutor::ServiceRunScriptProcess()
{
	if (!PathFileExists(m_szScriptPath)) {
		GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "【Error】脚本路径无效:%s", m_szScriptPath);
		return -1;
	}

	DWORD dwProcesses = 0;
	BOOL bResult = FALSE;
	DWORD dwRet = 0;
	DWORD dwSid = GetActiveSessionID();
	//DWORD dwSid2 = WTSGetActiveConsoleSessionId();

	HANDLE hProcess = NULL, hPToken = NULL, hUserTokenDup = NULL;
	if (!WTSQueryUserToken(dwSid, &hPToken)) {
		PROCESSENTRY32 procEntry;
		DWORD dwPid = 0;
		HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
		if (hSnap == INVALID_HANDLE_VALUE) {
			return FALSE;
		}

		procEntry.dwSize = sizeof(PROCESSENTRY32);
		if (Process32First(hSnap, &procEntry)) {
			do {
				if (_tcsicmp(procEntry.szExeFile, _T("explorer.exe")) == 0) {
					DWORD exeSessionId = 0;
					if (ProcessIdToSessionId(procEntry.th32ProcessID, &exeSessionId) && exeSessionId == dwSid) {
						dwPid = procEntry.th32ProcessID;
						break;
					}
				}

			} while (Process32Next(hSnap, &procEntry));
		}
		CloseHandle(hSnap);

		// explorer进程不存在
		if (dwPid == 0) {
			return FALSE;
		}

		hProcess = ::OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwPid);
		if (hProcess == NULL) {
			return FALSE;
		}

		if(!::OpenProcessToken(hProcess, TOKEN_ALL_ACCESS_P,&hPToken)) {
			CloseHandle(hProcess);
			return FALSE;
		}
	}

	if (hPToken == NULL)
		return -1;

	/*if ( !ImpersonateLoggedOnUser(hPToken) )
	{
		printf("模拟当前用户登录失败\n");
	}*/

	TOKEN_LINKED_TOKEN admin;
	bResult = GetTokenInformation(hPToken, (TOKEN_INFORMATION_CLASS)19, &admin, sizeof(TOKEN_LINKED_TOKEN), &dwRet);

	if (!bResult) {// vista 以前版本不支持TokenLinkedToken
		TOKEN_PRIVILEGES tp;
		LUID luid;
		if (LookupPrivilegeValue(NULL,SE_DEBUG_NAME,&luid)) {
			tp.PrivilegeCount =1;
			tp.Privileges[0].Luid =luid;
			tp.Privileges[0].Attributes =SE_PRIVILEGE_ENABLED;
		}
		DuplicateTokenEx(hPToken,MAXIMUM_ALLOWED,NULL,SecurityIdentification,TokenPrimary,&hUserTokenDup);
	} else {
		hUserTokenDup = admin.LinkedToken;
	}

	LPVOID pEnv =NULL;
	DWORD dwCreationFlags = CREATE_PRESERVE_CODE_AUTHZ_LEVEL;

	if(CreateEnvironmentBlock(&pEnv,hUserTokenDup,TRUE))
	{
		dwCreationFlags|=CREATE_UNICODE_ENVIRONMENT;
	}
	else
	{
		pEnv = NULL;
	}

	ZeroMemory( &m_si, sizeof(m_si) );
	m_si.cb = sizeof(m_si);
	m_si.dwFlags = STARTF_USESHOWWINDOW;
	m_si.wShowWindow = SW_HIDE;
	ZeroMemory( &m_pi, sizeof(m_pi) );

	// 重定向输出;
	RedirectSubprocessStdout(&m_si);
	// 恢复日志线程;
	ResumeThread(m_hStdoutLogThread);
	ResumeThread(m_hStderrLogThread);

	// 强制stdion, stdout和stderr完全无缓冲:python -u
	TCHAR szCommandLine[MAX_PATH] = { 0 };
	if (_tcslen(m_szExtraSentence))
		_stprintf_s(szCommandLine, _T("python -W ignore -u \"%s\" \"%s\""), m_szScriptPath, m_szExtraSentence);
	else
		_stprintf_s(szCommandLine, _T("python -W ignore -u \"%s\""), m_szScriptPath);

	bResult = CreateProcessAsUser(
		hUserTokenDup,                     // client's access token
		NULL,    // file to execute
		szCommandLine,                 // command line
		NULL,            // pointer to process SECURITY_ATTRIBUTES
		NULL,               // pointer to thread SECURITY_ATTRIBUTES
		TRUE,              // handles are not inheritable
		dwCreationFlags,     // creation flags
		pEnv,               // pointer to new environment block
		NULL,               // name of current directory
		&m_si,               // pointer to STARTUPINFO structure
		&m_pi                // receives information about new process
		);	
	if ( !bResult )
		return -1;

	// 记录进程ID;
	if ( m_pCaseObj )
		m_pCaseObj->__dwPythonPID = m_pi.dwProcessId;

	GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "Python进程ID=%ld", m_pi.dwProcessId);
	// 等待进程完成退出.
	WaitForSingleObject(m_pi.hProcess, INFINITE);
	
	// 必须等待，等待管道输出结果;
	Sleep(6000);
	GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "脚本进程结束（%ld）",  m_pi.dwProcessId);
	// 关闭进程句柄. 
	CloseHandle(m_pi.hProcess);
	CloseHandle(m_pi.hThread);

	// 关闭管道前,先取消管道的IO操作;
	CancelIo(m_hStdOutWrite);
	CancelIo(m_hStdErrorWrite);
	// 同时需要关闭输出的管道,否则ReadFile会阻塞;
	CloseHandle(m_hStdOutWrite);
	m_hStdOutWrite = NULL;
	CloseHandle(m_hStdErrorWrite);
	m_hStdErrorWrite = NULL;

	// 重置;
	memset(&m_si, 0, sizeof(m_si));
	memset(&m_pi, 0, sizeof(m_pi));
	m_si.cb = sizeof(m_si);

	if (hUserTokenDup != NULL)
		CloseHandle(hUserTokenDup);
	if (hProcess != NULL)
		CloseHandle(hProcess);
	if (hPToken != NULL)
		CloseHandle(hPToken);
	if (pEnv != NULL)
		DestroyEnvironmentBlock(pEnv);

	return TRUE;
}

bool CPythonExecutor::StartThread()
{
	// 脚本进程运行线程;
	m_hWorkThread = CreateThread(NULL, 0, _WorkerThread, this, 0, NULL);
	if (!m_hWorkThread) {
		printf("Error:创建执行线程失败\n");
		return false;
	}

	// 标准输出重定向线程;
	m_hStdoutLogThread = CreateThread(NULL, 0, _StdoutLogExportThread, this, CREATE_SUSPENDED, NULL);
	if (!m_hStdoutLogThread) {
		printf("Error:创建日志线程失败\n");
		return false;
	}

	// 标准错误重定向线程;
	m_hStderrLogThread = CreateThread(NULL, 0, _StderrLogExportThread, this, CREATE_SUSPENDED, NULL);
	if (!m_hStderrLogThread) {
		printf("Error:创建日志线程失败\n");
		return false;
	}

	return true;
}

void CPythonExecutor::EndWorkThread()
{
	// 结束进程;
	EndSubprocess();
	// 等待3秒,是否能自主结束线程;
	if (m_hWorkThread) {
		if (WaitForSingleObject(m_hWorkThread, 3000) == WAIT_OBJECT_0) {
			CloseHandle(m_hWorkThread);
			m_hWorkThread = NULL;
			return;
		}
	}

	// 手动结束线程;
	if ( m_hWorkThread ) {
		// 尝试5次结束行为;
		for (int i = 0; i < 5; i++) {
			if (TerminateThread(m_hWorkThread, 0))
				break;
		}

		CloseHandle(m_hWorkThread);
		m_hWorkThread = NULL;
	}
}

void CPythonExecutor::EndLogThread()
{
	// 标记结束;
	m_bStopLogExport = TRUE;
	Sleep(20); // 等待关闭线程完成一次循环并退出;
	// 同时关闭重定向的句柄;
	EndSubprocessStdOut();
	// 等待5秒,是否能自主结束线程;
	if ( m_hStdoutLogThread ) {
		if ( WaitForSingleObject(m_hStdoutLogThread, 5000) == WAIT_OBJECT_0 ) {
			CloseHandle(m_hStdoutLogThread);
			m_hStdoutLogThread = NULL;
		}	
	}

	// 手动结束线程;
	if (m_hStdoutLogThread) {
		TerminateThread(m_hStdoutLogThread, 0);
		// 异步结束,需要等待(5秒)系统完成操作;
		if ( WaitForSingleObject(m_hStdoutLogThread, 5000) == WAIT_OBJECT_0 ) {
			CloseHandle(m_hStdoutLogThread);
			m_hStdoutLogThread = NULL;
		}
		else
		{
			GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, _T("结束标准输出日志线程失败"));
		}
	}

	//////////////////////////////////////////////////////////////////////////
	if ( m_hStderrLogThread ) {
		if ( WaitForSingleObject(m_hStderrLogThread, 5000) == WAIT_OBJECT_0 ) {
			CloseHandle(m_hStderrLogThread);
			m_hStderrLogThread = NULL;
		}	
	}

	// 手动结束线程;
	if (m_hStderrLogThread) {
		TerminateThread(m_hStderrLogThread, 0);
		// 异步结束,需要等待(5秒)系统完成操作;
		if ( WaitForSingleObject(m_hStderrLogThread, 5000) == WAIT_OBJECT_0 ) {
			CloseHandle(m_hStderrLogThread);
			m_hStderrLogThread = NULL;
		}	
		else
		{
			GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, _T("结束标准错误日志线程失败"));
		}
	}
}

void CPythonExecutor::EndThread()
{
	// 先关闭日志线程;
	EndLogThread();
	// 再关闭脚本线程;
	EndWorkThread();
}

BOOL CPythonExecutor::EndSubprocess()
{
	// 注意TerminateProcess是异步的;
	TerminateProcess(m_pi.hProcess, 0);
	// 异步结束,需要等待(5秒)系统完成操作;
	if ( WaitForSingleObject(m_pi.hProcess, 5000) == WAIT_OBJECT_0 ) {
		GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, _T("结束进程（%ld）成功"), m_pi.dwProcessId);
		return true;
	}

	GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, _T("结束进程（%ld）失败"), m_pi.dwProcessId);
	return false;
}

void CPythonExecutor::EndSubprocessStdOut()
{
	OutputDebugString("关闭重定向子进程标准输出管道句柄\n");
	// 关闭重定向句柄;
	if (m_hStdErrorRead)
		CloseHandle(m_hStdErrorRead);
	m_hStdErrorRead = NULL;

	if (m_hStdErrorWrite)
		CloseHandle(m_hStdErrorWrite);
	m_hStdErrorWrite = NULL;

	// 只有子进程方式才关闭句柄;
	if ( m_hStdOutWrite)
		CloseHandle(m_hStdOutWrite);
	m_hStdOutWrite = NULL;

	if (m_hStdOutRead)
		CloseHandle(m_hStdOutRead);
	m_hStdOutRead = NULL;	
}

DWORD CPythonExecutor::GetActiveSessionID()
{
	DWORD dwSessionId = 0;
	PWTS_SESSION_INFO pSessionInfo = NULL;
	DWORD dwCount = 0;

	WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE, 0, 1, &pSessionInfo, &dwCount);

	for(DWORD i = 0; i < dwCount; i++) {
		WTS_SESSION_INFO si = pSessionInfo[i];
		if(WTSActive == si.State) {
			dwSessionId = si.SessionId;
			break;
		}
	}

	WTSFreeMemory(pSessionInfo);

	return dwSessionId;
}

bool CPythonExecutor::InitScript(std::string strScript, std::string strLogPath, std::string strScriptCmd)
{
	// 判断脚本是否存在;
	if (!PathFileExists(strScript.c_str())) {
		GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "【Error】脚本文件不存在:%s\n", strScript.c_str());
		return false;
	}

	// 判断日志文件路径是否可创建;
	if (!PathFileExists(strLogPath.c_str())) {
		// 创建路径;
		if (!GLOBAL::MKDIR(strLogPath.c_str())) {
			GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "【Error】创建目录失败：%s\n", strLogPath.c_str());
			return false;
		}

		// 创建文件;
		std::ofstream flog(strLogPath.c_str());
		if ( flog.bad() ) {
			GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, "【Error】创建文件失败：%ld,%s\n", GetLastError(), strLogPath.c_str());
			return false;
		}
		flog.close();
	}

	// 赋值参数;
	_tcscpy_s(m_szScriptPath, strScript.c_str());
	_tcscpy_s(m_szLogPath, strLogPath.c_str());
	_tcscpy_s(m_szExtraSentence, strScriptCmd.c_str());

	return true;
}

bool CPythonExecutor::StartScript()
{
	if (StartThread()) {
		Sleep(100);
		return true;
	}

	return false;
}

void CPythonExecutor::StopScript()
{
	// 如果是子进程运行脚本,停止线程时kill进程;
	EndSubprocess();
	// 结束线程;
	EndWorkThread();
}

bool CPythonExecutor::IsScriptOver()
{
	if ( WaitForSingleObject(m_hStdoutLogThread, 0) == WAIT_OBJECT_0  && 
		 WaitForSingleObject(m_hStderrLogThread, 0) == WAIT_OBJECT_0 && 
		 WaitForSingleObject(m_hWorkThread, 0) == WAIT_OBJECT_0)
	{
		// 标记任务完成;
		m_pCaseObj->_nExecutionState = SATHTTP::EXECUTED;
		// 如果任务结果为空,标记成功;
		if ( m_pCaseObj->_nExecutionResult == SATHTTP::NONE )
			m_pCaseObj->_nExecutionResult = SATHTTP::SUCCESS;
		GLOBAL::WriteTextLog(GLOBAL::SAT_PYE, _T("====>脚本(%s)已完成"), m_pCaseObj->strCaseName.c_str());
		return true;
	}

	return false;
}