#include "stdafx.h"

ULONG_PTR RvaToOffset(IMAGE_NT_HEADERS *, ULONG_PTR);
BOOL GetExport(BYTE *, ULONG_PTR *, WORD *, CHAR *); 

BOOL GetNtoskrnlOriginalBytes(PRESTORE_INFO pRestoreInfo)
{
	TCHAR Buffer[MAX_PATH];

	GetSystemDirectory(Buffer, MAX_PATH);

	_tcscat(Buffer, _T("\\ntoskrnl.exe"));

	HANDLE hFile = CreateFile(Buffer, GENERIC_READ, FILE_SHARE_READ,
		NULL, OPEN_EXISTING, 0, NULL);

	if (hFile == INVALID_HANDLE_VALUE)
		return FALSE;

	DWORD FileSize = GetFileSize(hFile, NULL);

	BYTE *ptrNtoskrnl = (BYTE *) VirtualAlloc(NULL, FileSize, 
		MEM_COMMIT, PAGE_READWRITE);

	if (ptrNtoskrnl == NULL)
	{
		CloseHandle(hFile);
		return FALSE;
	}

	DWORD BR;

	if (!ReadFile(hFile, ptrNtoskrnl, FileSize, &BR, NULL))
	{
		VirtualFree(ptrNtoskrnl, 0, MEM_RELEASE);
		CloseHandle(hFile);
		return FALSE;
	}

	CloseHandle(hFile);

	IMAGE_DOS_HEADER *ImgDosHdr = (IMAGE_DOS_HEADER *) ptrNtoskrnl;

	IMAGE_NT_HEADERS *ImgNtHdrs = (IMAGE_NT_HEADERS *)
		&ptrNtoskrnl[ImgDosHdr->e_lfanew];

	ULONG_PTR EP_Rva = 0;

	if (!GetExport(ptrNtoskrnl, &EP_Rva, NULL, "KeAttachProcess"))
	{
		VirtualFree(ptrNtoskrnl, 0, MEM_RELEASE);
		return FALSE;
	}

	BYTE *ptr = (BYTE *) (EP_Rva + (ULONG_PTR) ptrNtoskrnl);

	memcpy(pRestoreInfo->KeAttachProcessPatch, ptr, 5);

	VirtualFree(ptrNtoskrnl, 0, MEM_RELEASE);

	return TRUE;
}

BOOL CollectInformation(PRESTORE_INFO pRestoreInfo)
{
	if (SDT_GetOriginalFunctions(pRestoreInfo) == FALSE)
		return FALSE;

	return GetNtoskrnlOriginalBytes(pRestoreInfo);
}