#include "StdAfx.h"
#include "GuardMain.h"
#include <ProfInfo.h>
#include <io.h>
#include <tlhelp32.h>
#include "Config.h"


SINGLETON_IMPLEMENT(CGuardMain)


CGuardMain::CGuardMain(void) : m_hFileConfig(NULL), m_hMapFileConfig(NULL), m_pThreadWatch(NULL), m_pThreadTimer(NULL), m_pProcessInfo(NULL)
{
}


CGuardMain::~CGuardMain(void)
{
}

/*****************************************************************
**【函数名称】	__runWithOs
**【函数功能】	随系统启动
**【参数】		
**【返回值】	
****************************************************************/
void CGuardMain::__runWithOs( void )
{
	//获得当前应用程序路径
	TCHAR path[MAX_PATH] = { 0 };
	if( !GetModuleFileName( NULL, path, MAX_PATH ) )
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_WARNING, _T("{Guard}: 设置随系统启动失败, 无法获取系统路径"));
		return;
	}

	CRegKey key;
	if (key.Open(HKEY_LOCAL_MACHINE, REG_STARTUP_PATH, KEY_WRITE | KEY_READ) != ERROR_SUCCESS)
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_WARNING, _T("{Guard}: 设置随系统启动失败, 无法打开系统接口"));
		return;
	}

	CString Path2Write;
	Path2Write.Format(_T("\"%s\""), path);
	if(key.SetStringValue(REG_STARTUP_ITEM, Path2Write, REG_SZ) != ERROR_SUCCESS)
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_WARNING, _T("{Guard}: 设置随系统启动失败, 无法建立启动项"));
}

/*****************************************************************
**【函数名称】	__loadConfigInfo
**【函数功能】	加载配置信息
**【参数】		
**【返回值】	BOOL
****************************************************************/
bool CGuardMain::__loadConfigInfo( void )
{
	// 加载配置文件
	if(_access(FILE_PATH, 0) == -1) 
		CreateDirectoryA(FILE_PATH, NULL);
	m_hFileConfig = CreateFile(FILE_NAME, 
		GENERIC_READ|GENERIC_WRITE, 
		FILE_SHARE_READ|FILE_SHARE_WRITE, 
		NULL, 
		OPEN_ALWAYS, 
		FILE_ATTRIBUTE_NORMAL, 
		NULL);
	if(m_hFileConfig == NULL) 
		return false;


	// 打开内存映射文件
	m_hMapFileConfig = CreateFileMapping(m_hFileConfig, NULL, PAGE_READWRITE, 0, sizeof(ProcessInfo)*WATCH_MAX_COUNT, VIEW_NAME);
	if(m_hMapFileConfig == NULL)
	{
		CloseHandle(m_hFileConfig);
		m_hFileConfig = NULL;
		return false;
	}

	// 映射文件对象到进程地址空间
	LPVOID hViewFile = MapViewOfFile(m_hMapFileConfig, FILE_MAP_ALL_ACCESS, 0, 0, sizeof(ProcessInfo)*WATCH_MAX_COUNT);
	if(hViewFile == NULL)
	{
		CloseHandle(m_hMapFileConfig);
		m_hMapFileConfig = NULL;
		CloseHandle(m_hFileConfig);
		m_hFileConfig = NULL;
		return false;
	}

	// 保存内存指针
	if(m_hFileConfig != NULL) 
	{
		CloseHandle(m_hFileConfig);
		m_hFileConfig = NULL;
	}
	if(m_hMapFileConfig != NULL) 
	{
		CloseHandle(m_hMapFileConfig);
		m_hMapFileConfig = NULL;
	}
	m_pProcessInfo = (ProcessInfo*)hViewFile;

	return true;
}

/*****************************************************************
**【函数名称】	__onWatch
**【函数功能】	处理监控函数
**【参数】		
**【返回值】	
****************************************************************/
void CGuardMain::__onWatch( void )
{
	for(int i = 0; i < WATCH_MAX_COUNT; i++)
	{
		// 屏蔽路径为空的项
		if(strcmp(m_pProcessInfo[i].m_szName, _T("")) == 0 ||
			strcmp(m_pProcessInfo[i].m_szPath, _T("")) == 0)
			continue;

		// 从系统中查找进程
		DWORD dwProcessId = 0;
		if(!__searchProcess(m_pProcessInfo[i].m_szName, dwProcessId))
		{
			// 搜索不到进程，启动进程
 			if(m_pProcessInfo[i].m_nErrorCount >= 2)
 			{
				m_pProcessInfo[i].m_nErrorCount = 0;
				ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_WARNING, _T("{Guard}: %s监控失败, 将启动该进程, path = %s"), 
					m_pProcessInfo[i].m_szName, m_pProcessInfo[i].m_szPath);
				__runProcess(m_pProcessInfo[i]);
			}
			else
			{
				m_pProcessInfo[i].m_nErrorCount++;
			}
		}
		else
		{
			if(m_pProcessInfo[i].m_bUseActiveFlag) 
			{
				// 进程存在，判断进程是否假死
				if(m_pProcessInfo[i].m_bActiveFlag)
				{
					m_pProcessInfo[i].m_nErrorCount = 0;
					m_pProcessInfo[i].m_bActiveFlag = FALSE;
				}
				else
				{
					if(m_pProcessInfo[i].m_nErrorCount >= 2)
					{
						m_pProcessInfo[i].m_nErrorCount = 0;

						// 杀死进程
						ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_WARNING, _T("{Guard}: 检测不到%s心跳, 将关闭该进程"), 
							m_pProcessInfo[i].m_szName);
						__killProcess(dwProcessId);				
					}
					else
					{
						m_pProcessInfo[i].m_nErrorCount++;
					}
				}
			}
			else
			{
				// 不启用心跳通讯，检测到进程存在
				m_pProcessInfo[i].m_nErrorCount = 0;
			}
		} // end if
	} // end for
}

/*****************************************************************
**【函数名称】	__onTimer
**【函数功能】	处理定时器函数
**【参数】		
**【返回值】	
****************************************************************/
void CGuardMain::__onTimer( void )
{
	CTime TimeCurr = CTime::GetCurrentTime();
	CString strTime = TimeCurr.Format(_T("%H"));

	for(int i = 0; i < WATCH_MAX_COUNT; i++)
	{
		if(strcmp(m_pProcessInfo[i].m_szName, _T("")) == 0) continue;

		// 定时器时长增加
		m_pProcessInfo[i].m_nRunTimer++;

		// 运行超过指定时长，并且当前为01:00，重启该进程
		if(m_pProcessInfo[i].m_nOnTimerReStart != 0 && m_pProcessInfo[i].m_nRunTimer >= m_pProcessInfo[i].m_nOnTimerReStart 
			&& strTime == CConfig::processRestartOclock())
		{
			ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_NORMAL, _T("{Guard}: %s已运行%d小时, 达到定时重启时限(%d小时), 将重新启动"), m_pProcessInfo[i].m_szName, 
				m_pProcessInfo[i].m_nRunTimer, m_pProcessInfo[i].m_nOnTimerReStart);

			// 重启进程
			DWORD dwProcessId = 0;
			if(__searchProcess(m_pProcessInfo[i].m_szName, dwProcessId))
			{
				// 杀死进程
				__killProcess(dwProcessId);
			}

			// 重启进程
			__runProcess(m_pProcessInfo[i]);		
		}
	} // endfor
}

/*****************************************************************
**【函数名称】		__getLastError
**【函数功能】      获取错误信息
**【参数】         
**【返回值】
****************************************************************/
CString CGuardMain::__getLastError( void )
{
	LPVOID lpMsgBuf;
	FormatMessage(
		FORMAT_MESSAGE_ALLOCATE_BUFFER |
		FORMAT_MESSAGE_FROM_SYSTEM |
		FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL,
		GetLastError(),
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), // Default language
		(LPTSTR) &lpMsgBuf,
		0,
		NULL
		);

	CString strErrMsg = _T("");
	strErrMsg.Format("%s", (LPCTSTR)lpMsgBuf);
	LocalFree(lpMsgBuf);

	return strErrMsg;
}

/*****************************************************************
**【函数名称】		__getTokenByName
**【函数功能】      得到登录用户名和口令
**【参数】          HANDLE	标识符句柄
					lpName	进程名称
**【返回值】
****************************************************************/
bool CGuardMain::__getTokenByName(HANDLE &hToken, CString strName)
{
	if(strName.IsEmpty())   
		return false;

    HANDLE hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if(hProcessSnap == INVALID_HANDLE_VALUE) 
		return false;     
      
	BOOL bRet = FALSE;     
	PROCESSENTRY32 pe32 = {0};     
    pe32.dwSize = sizeof(PROCESSENTRY32);     
    if(Process32First(hProcessSnap, &pe32))     
	{       
		do     
        {   
			if(strName.CompareNoCase(pe32.szExeFile) == 0)
			{   
				HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pe32.th32ProcessID);   
				bRet = OpenProcessToken(hProcess, TOKEN_ALL_ACCESS, &hToken);   
				CloseHandle(hProcessSnap);     
				return bRet != FALSE;   
			}   
		}     
        while(Process32Next(hProcessSnap, &pe32));
	}
      
    CloseHandle(hProcessSnap);     
    return  false;   
}

/*****************************************************************
**【函数名称】	__searchProcess
**【函数功能】	搜索进程
**【参数】		lpProcessName 进程名称
				dwProcessId 返回的进程ID
**【返回值】	
****************************************************************/
bool CGuardMain::__searchProcess(LPCTSTR lpProcessName, DWORD& dwProcessId)
{
	HANDLE hSnapshot;
	hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);	
	if(hSnapshot == INVALID_HANDLE_VALUE) 
		return false;

	bool bRet = false;
	PROCESSENTRY32 lppe;
	lppe.dwSize = sizeof(PROCESSENTRY32);
	if(Process32First(hSnapshot, &lppe)) 
	{
		do
		{
			if(strcmp(lppe.szExeFile, lpProcessName) == 0) 
			{
				dwProcessId = lppe.th32ProcessID;
				bRet = true;
				break;
			}
		}
	    while(Process32Next(hSnapshot, &lppe)); // 查找下一个进程
	}
   
   CloseHandle(hSnapshot);
   return bRet;
}

/*****************************************************************
**【函数名称】	__killProcess
**【函数功能】	杀死进程
**【参数】		dwProcessId 进程ID
**【返回值】	
****************************************************************/
void CGuardMain::__killProcess(DWORD dwProcessId)
{
	DWORD xCode;
    HANDLE hProc = OpenProcess(PROCESS_TERMINATE, false, dwProcessId);
    if(hProc && !GetExitCodeProcess(hProc, &xCode)) 
	{	
		// 结束进程
        TerminateProcess(hProc, xCode);
		CloseHandle(hProc);
		Sleep(3);
    }
}

/*****************************************************************
**【函数名称】		__runProcess
**【函数功能】      运行指定进程并前台显示
**【参数】          ProInfo 进程信息
**【返回值】
****************************************************************/
void CGuardMain::__runProcess(ProcessInfo& ProInfo)
{
	// 获取应用程序目录，设置子进程的当前目录
	CString strDircet = ProInfo.m_szPath;
	int nIndex = strDircet.ReverseFind('\\');
	if(-1 == nIndex)
	{
		nIndex = strDircet.ReverseFind('/');
	} 
	strDircet = strDircet.Left(nIndex+1);

	STARTUPINFO si;
	PROCESS_INFORMATION pi;

	ZeroMemory(&si, sizeof(si));
	si.cb = sizeof(si);
	ZeroMemory(&pi, sizeof(pi));

	BOOL bRet = CreateProcess(ProInfo.m_szPath,		// No module name (use command line)
		NULL,				// Command line
		NULL,				// Process handle not inheritable
		NULL,				// Thread handle not inheritable
		FALSE,			    // Set handle inheritance to FALSE
		0,					// No creation flags
		NULL,				// Use parent's environment block
		strDircet,			// Use parent's starting directory 
		&si,				// Pointer to STARTUPINFO structure
		&pi); 				// Pointer to PROCESS_INFORMATION structure

	// 运行时长清零
	if(bRet)
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_NORMAL, _T("{Guard}: 进程%s启动成功, Path = %s"), ProInfo.m_szName, ProInfo.m_szPath);
		ProInfo.m_nRunTimer = 0;
		CloseHandle( pi.hProcess );
		CloseHandle( pi.hThread );

		Sleep(5000);

	}
	else
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_WARNING, _T("{Guard}: 进程%s启动失败, Path = %s, Cause = %s"), 
			ProInfo.m_szName, ProInfo.m_szPath, __getLastError());
	}
}

/*****************************************************************
**【函数名称】		__runProcess
**【函数功能】      运行指定进程并前台显示
**【参数】          ProInfo 进程信息
**【返回值】
****************************************************************/
void CGuardMain::__runProcessByService(ProcessInfo& ProInfo)
{
	// 获取登录用户信息
	HANDLE	hToken = NULL;			
	if(!__getTokenByName(hToken, _T("EXPLORER.EXE")))
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_WARNING, _T("{Guard}: 启动进程%s时获取用户令牌失败, Cause = %s"), __getLastError());
		return; 
	}

	typedef BOOL (STDMETHODCALLTYPE FAR * LPFNLOADUSERPROFILE)(HANDLE hToken, LPPROFILEINFO lpProfileInfo);
	HMODULE hLib = LoadLibrary(_T("userenv.dll"));
	LPFNLOADUSERPROFILE LoadUserProfile = NULL;

#ifdef UNICODE
	LoadUserProfile = (LPFNLOADUSERPROFILE)GetProcAddress(hLib, _T("LoadUserProfileW"));
#else
	LoadUserProfile = (LPFNLOADUSERPROFILE)GetProcAddress(hLib, _T("LoadUserProfileA"));
#endif

	if(LoadUserProfile == NULL) return;
	PROFILEINFO profile = {0};
	profile.dwSize = sizeof(PROFILEINFO);
	profile.lpUserName = "administrator";
	if(LoadUserProfile(hToken, &profile) == NULL)  
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_WARNING, _T("{Guard}: 启动进程%s时加载用户配置文件失败, Cause = %s"), __getLastError());
		return;
	}

	PROCESS_INFORMATION pi;				// 进程信息 
	STARTUPINFO lp = {0};				// 新进程窗口信息
	lp.cb = sizeof(lp);
	lp.dwFlags = STARTF_USESHOWWINDOW|STARTF_USESTDHANDLES;
	lp.wShowWindow = SW_HIDE;

	// 启动新进程
	BOOL bRet = CreateProcessAsUser(hToken,
		ProInfo.m_szPath,
		NULL,
		NULL,
		NULL,
		FALSE,
		NORMAL_PRIORITY_CLASS,
		NULL, 
		NULL,
		&lp,
		&pi);

	CloseHandle(hToken);
}

/*****************************************************************
**【函数名称】	__watchThreadFun
**【函数功能】	监控线程类
**【参数】		
**【返回值】	
****************************************************************/
UINT CGuardMain::__watchThreadFun(LPVOID Param)
{
	CGuardMain* pGuard = (CGuardMain*)Param;
	ASSERT(pGuard != NULL);

	while(TRUE)
	{
		Sleep(WATCH_TIMER_SERVICE);
		pGuard->__onWatch();
	}
	return 0;
}

/*****************************************************************
**【函数名称】	__timerThreadFun
**【函数功能】	定时器线程函数
**【参数】		
**【返回值】	
****************************************************************/
UINT CGuardMain::__timerThreadFun(LPVOID Param)
{
	CGuardMain* pGuard = (CGuardMain*)Param;
	ASSERT(pGuard != NULL);

	while(TRUE)
	{
		// 每个小时观察一次
		Sleep(1000 * 3600);
		pGuard->__onTimer();
	}
	return 0;
}

/*****************************************************************
**【函数名称】	init
**【函数功能】	初始化监控控制类
**【参数】		
**【返回值】	
****************************************************************/
bool CGuardMain::init( void )
{
	//__runWithOs();

	// 加载配置信息
	if(!__loadConfigInfo()) 
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_ERROR, _T("{Guard}: 读取映射文件失败, Cause = %s"), __getLastError());
		return false;
	}

	// 创建监控线程
	m_pThreadWatch = AfxBeginThread(__watchThreadFun, this, 0, 0, CREATE_SUSPENDED, NULL);
	if(m_pThreadWatch == NULL) 
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_ERROR, _T("{Guard}: 启动守护监控线程失败"));
		return false;
	}

	// 创建定时器线程
	m_pThreadTimer = AfxBeginThread(__timerThreadFun, this, 0, 0, CREATE_SUSPENDED, NULL);
	if(m_pThreadTimer == NULL)
	{
		ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_ERROR, _T("{Guard}: 启动定时守护线程失败"));
		return false;
	}

	// 启动监控服务
	startWatchServer();

	ILogger::getInstance().log(LOG_CLASS_BUSI, LOG_LEVEL_NORMAL, _T("{Guard}: 守护功能启动成功"));
	return TRUE;
}

/*****************************************************************
**【函数名称】	startWatch
**【函数功能】	开始监控
**【参数】		
**【返回值】	
****************************************************************/
void CGuardMain::startWatchServer()
{
	if(m_pThreadWatch != NULL) m_pThreadWatch->ResumeThread();
	if(m_pThreadTimer != NULL) m_pThreadTimer->ResumeThread();
}

/*****************************************************************
**【函数名称】	stopWatch
**【函数功能】	停止监控
**【参数】		
**【返回值】	
****************************************************************/
void CGuardMain::stopWatchServer()
{
	if(m_pThreadWatch != NULL) m_pThreadWatch->SuspendThread();
	if(m_pThreadTimer != NULL) m_pThreadTimer->SuspendThread();
}