#include <windows.h>
#include <tchar.h>
#include <stdio.h>	// sprintf

#include "ServiceStatus.h"
#include "EventLog.h"
#include "IOCP.h"
#include "security.h"

#include "events.h"
#include "..\inc\service_iface.h"
#include "..\inc\kernel_iface.h"
#include "test.h"
#include "tests_w2k.h"
#include "..\inc\config.h"


CServiceStatus pfSrvStatus;
CEventLog pfEventLog (pfAppName);
CIOCP* piocp;

PTEST baseline[pfNTESTS], currTest[pfNTESTS];

// The completion port wakes for 1 of 2 reasons:
enum COMPKEY { 
   CK_SERVICECONTROL,   // A service control code
   CK_PIPE              // A client connects to our pipe
};

HANDLE hPipeCmd, hPipeData;

void logDebug (PTSTR msg) {
	
	PTSTR ppszStrings[] = {msg, NULL};


	pfEventLog.ReportEvent (EVENTLOG_INFORMATION_TYPE,
		0, MSG_DEBUG, CEventLog::REUSER_NOTAPPLICABLE, 1, (LPCTSTR*)ppszStrings, 0, NULL);
	
}

void logError (PTSTR err) {
	PTSTR ppszStrings[] = {err, NULL};

	pfEventLog.ReportEvent (EVENTLOG_ERROR_TYPE,
		0, MSG_ERROR, CEventLog::REUSER_NOTAPPLICABLE, 1, (LPCTSTR*)ppszStrings, 0, NULL);
}

void cleanup () {

	CloseHandle (hPipeCmd);
	CloseHandle (hPipeData);
//	pfEventLog.Uninstall();
}




DWORD WINAPI pfHandlerEx (DWORD dwControl, DWORD dwEventType,
						  PVOID pvEventData, PVOID pvContext) {

	DWORD dwReturn = ERROR_CALL_NOT_IMPLEMENTED;
	switch (dwControl) {
	case SERVICE_CONTROL_INTERROGATE:
		pfSrvStatus.ReportStatus();
		break;
	
#ifdef ALLOW_STOP_CONTROL
	case SERVICE_CONTROL_STOP:
#endif
	case SERVICE_CONTROL_SHUTDOWN:
		cleanup();
		pfSrvStatus.SetUltimateState(SERVICE_STOPPED, 500);	
		piocp->PostStatus (CK_SERVICECONTROL, dwControl);
		dwReturn = NO_ERROR;
		break;
	}

	return dwReturn;
}

BOOL pfmodule_present() {
	if (pfSrv (PF_QUERY, 0) == PF_ANSWER) return TRUE;
	else return FALSE;
}

BOOL pfsrv_restarted() {
	return pfSrv (PF_SERVICE_STARTED, 0) ;
}


void WINAPI pfServiceMain (DWORD dwArgc, PTSTR* pszArgv) {
	
	pfSrvStatus.Initialize (pfSrvName, pfHandlerEx, 0, FALSE);
#ifdef ALLOW_STOP_CONTROL
	pfSrvStatus.AcceptControls ( SERVICE_CONTROL_STOP );
#endif
	pfSrvStatus.AcceptControls ( 0 );

	pfSrvStatus.AdvanceState (50*1000, 1);

	
	pfEventLog.Install (
		EVENTLOG_INFORMATION_TYPE | 
		EVENTLOG_WARNING_TYPE | 
		EVENTLOG_ERROR_TYPE,
		NULL, NULL, 1, NULL);

	
	if (!pfmodule_present()) {
		logError ("pfDriver not loaded, exiting...");
		exit (1);
	}
#ifdef PROTECT_SERVICE_RESTART
	if (pfsrv_restarted()) {
		logError ("pfServie has been restarted, someone is cheating you!!!");
		exit (1);
	}
#endif

	logDebug ("going sleep...");

	pfSrvStatus.ReportUltimateState();

	// wait for the system to boot
	// take measurments when system is in clear state
	Sleep (INITIAL_SLEEP_TIME);
	
	char msg[1024];
	sprintf (msg, "generating system baseline (niter = %d)...",
		BASELINE_NITER);
	logDebug (msg);
	int ntests = 0;
	int niter = BASELINE_NITER;

	
	#define _ADDTEST(tClass) {				\
		baseline[ntests] = new tClass;		\
		currTest[ntests] = new tClass;		\
		baseline[ntests]->go (niter);		\
		ntests ++;							\
	}
	
	
	_ADDTEST(test_Processes)
	_ADDTEST(test_FindFile)
	_ADDTEST(test_OpenFile)
	_ADDTEST(test_null)
	_ADDTEST(test_recvNull)
	_ADDTEST(test_RegOpenKey)
	_ADDTEST(test_RegEnum)
	_ADDTEST(test_RegEnumNull)

	#undef _ADDTEST
	if (ntests != pfNTESTS) logError ("ntests != pfNTESTS");
	//logDebug ("baseline generation complete, creating pipe...");

	CIOCP iocp (0);
	piocp = &iocp;
	SECURITY_ATTRIBUTES sa;
	if (!InitializePipeSecurity (&sa))
		logError("Cannot initialize SD for the pipe");

	hPipeCmd = CreateNamedPipe (pfPipeCmd,
		PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED,
		PIPE_TYPE_BYTE, 1, sizeof (SRVCMD), sizeof (SRVCMD), 0, &sa);

	if (hPipeCmd == INVALID_HANDLE_VALUE) {
		logError ("Cannot create pipe, exiting...");
		exit (1);
	}

	iocp.AssociateDevice(hPipeCmd, CK_PIPE);
	
	// Pend an asynchronous connect against the pipe
	OVERLAPPED o, *po;
	ZeroMemory(&o, sizeof(o));

	ConnectNamedPipe(hPipeCmd, &o);


	
	/* now, we're ready to serve client's queries */

	char pfver[256];
	sprintf (pfver, "%s", PFSRV_VERSION);
	PTSTR ppszStrings[] = {pfver, NULL};


	pfEventLog.ReportEvent (EVENTLOG_INFORMATION_TYPE,
		0, MSG_START_OK, CEventLog::REUSER_NOTAPPLICABLE, 1,
		(LPCTSTR*)ppszStrings, 0, NULL);
	
	
	DWORD dwControl, dwNumBytes;
	SRVCMD srvcmd;
	ULONG_PTR CompKey;
	int tno;
	do {
		iocp.GetStatus(&CompKey, &dwNumBytes, &po);
		dwControl = dwNumBytes;
		switch (CompKey) {
		
		case CK_PIPE:
			if (hPipeCmd != 0) {
				ReadFile (hPipeCmd, (LPVOID)&srvcmd,
					sizeof(srvcmd), &dwNumBytes, NULL);
				
								
				if (dwNumBytes != sizeof(srvcmd)) {
					logError("read from pipe error!");
					break;
				}

				switch (srvcmd.cmd) {
				case SRVCMD_CHECK_MODULE:
					if (pfmodule_present())
						srvcmd.resp = SRVRESP_MODULE_PRESENT;
					else srvcmd.resp = SRVRESP_MODULE_NOT_PRESENT;
					break;
	
				case SRVCMD_RUN_TEST:
					tno = srvcmd.testno;
					currTest[tno]->go (srvcmd.niter);
					
					srvcmd.peekClear = baseline[tno]->getPeek();
					srvcmd.peekPrClear = baseline[tno]->getPeekPr();
					
					srvcmd.peekCurr = currTest[tno]->getPeek();
					srvcmd.peekPrCurr = currTest[tno]->getPeekPr();

					strcpy (srvcmd.shortname, baseline[tno]->getNameShort());
					//strcpy (srvcmd.shortname, "AAAA");
					srvcmd.resp = SRVRESP_OK;
					break;
				default:
					srvcmd.resp = SRVRESP_INVALID_CMD;
				}
		
				WriteFile (hPipeCmd, (LPVOID)&srvcmd, sizeof(srvcmd),
					&dwNumBytes, NULL);
				if (dwNumBytes != sizeof(srvcmd)) {
					logError("write to pipe error!");
					break;
				}

				FlushFileBuffers(hPipeCmd);
				DisconnectNamedPipe(hPipeCmd);

				//Allow another client to connect 
				ZeroMemory(&o, sizeof(o));
				ConnectNamedPipe(hPipeCmd, &o);
			} else {
				// We get here when the pipe is closed
				logError ("ConnectNamedPipe returned 0!");
			}
			break;

		
		case CK_SERVICECONTROL:
	//		logDebug ("CK_SERVICECONTROL");
			switch (dwControl) {
				case SERVICE_CONTROL_STOP:
					logDebug ("SERVICE_CONTROL_STOP");
					pfSrvStatus.ReportUltimateState();
					break;
				case SERVICE_CONTROL_SHUTDOWN:
					logDebug ("SERVICE_CONTROL_SHUTDOWN");
					pfSrvStatus.ReportUltimateState();
					break;
			}
			break;
		} 
	} while (pfSrvStatus != SERVICE_STOPPED);
	return;
}



int main () {

	SERVICE_TABLE_ENTRY ServiceTable[] = {
		{pfSrvName, pfServiceMain},
		{NULL, NULL}
	};

	StartServiceCtrlDispatcher (ServiceTable);
	
	return 0;
}