/*
*	Module Name: 
*		WorkingSetWatch.cpp
*
*	Abstract:
*		Tracks page faults that occur within the process.
*		
*		NOTE: This is not compatible with Wow64 and must be run as a 64-bit
*		program on x64 and a 32-bit program on x86.
*
*	Author:
*		Nemanja (Nemi) Mulasmajic <nm@triplefault.io>
*			http://triplefault.io
*/
#pragma warning(disable: 4710)
#pragma comment(lib, "psapi.lib")

#pragma warning(push, 0)
#include <windows.h>
#include <stdio.h>
#include <Psapi.h>
#pragma warning(pop)

// A pseudo-handle that represents the current process.
#define NtCurrentProcess() ((HANDLE)-1)

// The size of an architecture page on x86/x64.
#define PAGE_SIZE ((SIZE_T)0x1000)

// Aligns a memory address (Va) on a page boundary.
#define PAGE_ALIGN(Va) ((PVOID)((ULONG_PTR)(Va) & ~(PAGE_SIZE - 1)))

// The initial size of the GetWsChanges/Ex array.
#define INITIAL_RECORD_SIZE 1000

/*
*	Discovers the owning process and its path from a given thread.
*/
bool GetProcessDataFromThread(_In_ DWORD ThreadId, _Out_ DWORD& ProcessId, _Out_writes_z_(MAX_PATH) wchar_t ProcessPath[MAX_PATH])
{
	bool status = false;
	
	ProcessId = 0;
	ProcessPath[0] = NULL;

	// Get a handle to the thread.
	HANDLE Thread = OpenThread(THREAD_QUERY_LIMITED_INFORMATION, FALSE, ThreadId);
	if (Thread)
	{
		// Extract the PID of the owning process.
		ProcessId = GetProcessIdOfThread(Thread);
		if (ProcessId)
		{
			// Get a handle to the process.
			HANDLE Process = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, ProcessId);
			if (Process)
			{
				// Extract the path of the process.
				status = (GetModuleFileNameExW(Process, NULL, ProcessPath, MAX_PATH) != 0);
			
				CloseHandle(Process);
			}
		}

		CloseHandle(Thread);
	}

	return status;
}

/*
*	The entry point.
*/
int main(void)
{
	int status = -1;

	PBYTE AllocatedBuffer = NULL;
	PPSAPI_WS_WATCH_INFORMATION_EX WatchInfoEx = NULL;

	const DWORD CurrentProcessId = GetCurrentProcessId();
	printf("[+] PID: %lu\n", CurrentProcessId);

#if defined(_M_IX86)
	// Can't run on Wow64 (32-bit on 64-bit OS).
	BOOL Wow64Process = FALSE;
	if (IsWow64Process(NtCurrentProcess(), &Wow64Process) && Wow64Process)
	{
		fprintf(stderr, "[-] ERROR: This process cannot be run under Wow64.\n");
		goto Cleanup;
	}
#endif

	// Initiate monitoring of the working set for this process.
	if (!InitializeProcessForWsWatch(NtCurrentProcess()))
	{
		fprintf(stderr, "[-] ERROR: Failed to initialize process for working set watch. InitializeProcessForWsWatch failed with error: %lu.\n", GetLastError());
		goto Cleanup;
	}

	// Allocate a buffer that we'll track and see when it's paged in to our 
	// process.
	// 
	// NOTE:
	// As long as we don't access this buffer directly, as an optimization, 
	// Windows will not map this into our process' working set.
	AllocatedBuffer = (PBYTE)VirtualAlloc(NULL, PAGE_SIZE, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
	if (!AllocatedBuffer)
	{
		fprintf(stderr, "[-] ERROR: Failed to allocate %Iu bytes for page faulting test buffer.\n", PAGE_SIZE);
		goto Cleanup;
	}

	printf("[+] Allocated buffer at 0x%p.\n", AllocatedBuffer);

	// This buffer will constantly increase in size if we're unable to fill the 
	// data with what we currently have.
	DWORD WatchInfoSize = (sizeof(PSAPI_WS_WATCH_INFORMATION_EX) * INITIAL_RECORD_SIZE);
	WatchInfoEx = (PPSAPI_WS_WATCH_INFORMATION_EX)malloc(WatchInfoSize);
	if (!WatchInfoEx)
	{
		fprintf(stderr, "[-] ERROR: Failed to allocate %lu bytes.\n", WatchInfoSize);
		goto Cleanup;
	}

	// Loop until we discover that our memory location is mapped in.
	while (TRUE)
	{
		// Each iteration of the loop we want to make sure that the watch array
		// (collection of pages mapped into our process' working set since the last
		// time we called the API) is reset.
		memset(WatchInfoEx, 0, WatchInfoSize);
	
		// Retrieve the newly mapped pages into our process' working set.
		if (!GetWsChangesEx(NtCurrentProcess(), WatchInfoEx, &WatchInfoSize))
		{
			DWORD ErrorCode = GetLastError();

			// This really isn't an error. This just means that no new pages
			// have been mapped into our process' VA since the last time
			// we called GetWsChangesEx.
			if (ErrorCode == ERROR_NO_MORE_ITEMS)
			{
				// Wait a little bit before trying again.
				Sleep(1);
				continue;
			}

			// Any other error code is bad.
			if (ErrorCode != ERROR_INSUFFICIENT_BUFFER)
			{
				fprintf(stderr, "[-] ERROR: GetWsChangesEx failed with error: %lu.\n", ErrorCode);
				goto Cleanup;
			}

			// If we get this far, we need to increase the buffer size. 
			WatchInfoSize *= 2;

			free(WatchInfoEx);
			WatchInfoEx = (PPSAPI_WS_WATCH_INFORMATION_EX)malloc(WatchInfoSize);

			if (!WatchInfoEx)
			{
				fprintf(stderr, "[-] ERROR: Failed to allocate %lu bytes.\n", WatchInfoSize);
				goto Cleanup;
			}

			continue;
		}

		// At this point, we've successfully returned an array of all the pages
		// that were mapped into our process' working set.

		// Let's check to see if we found the page we care about.
		bool bFound = false;
		for (size_t i = 0;; ++i)
		{
			PPSAPI_WS_WATCH_INFORMATION_EX info = &WatchInfoEx[i];

			// The array is terminated with a structure whose FaultingPc member is NULL.
			if (info->BasicInfo.FaultingPc == NULL)
				break;

			// The page fault may be on some offset on the page. We make 
			// sure to align this on a page boundary and then check to see if it's 
			// our allocated buffer (which should already be aligned on a page 
			// boundary). 
			PVOID FaultingPageVa = PAGE_ALIGN(info->BasicInfo.FaultingVa);
			if (FaultingPageVa == AllocatedBuffer)
			{
				printf("[+] 0x%p (0x%p) was mapped by 0x%p (TID: %lu).\n", FaultingPageVa, info->BasicInfo.FaultingVa, info->BasicInfo.FaultingPc, (DWORD)info->FaultingThreadId);
				
				DWORD ProcessId;
				wchar_t ProcessPath[MAX_PATH];

				// Get identifying information about the process that caused the page fault.
				if (GetProcessDataFromThread((DWORD)info->FaultingThreadId, ProcessId, ProcessPath))
					printf("\t--> %S (PID: %lu).\n", ProcessPath, ProcessId);

				// Signal to the outer loop that we're done.
				bFound = true;
				break;
			}
		}

		// Keep iterating until we've detected a fault.
		if (bFound)
			break;
	}

	status = 0;

Cleanup:
	// 'free' the 'malloc's.
	if (WatchInfoEx)
	{
		free(WatchInfoEx);
		WatchInfoEx = NULL;
	}

	if (AllocatedBuffer)
	{
		VirtualFree(AllocatedBuffer, 0, MEM_RELEASE);
		AllocatedBuffer = NULL;
	}

	// Wait for [ENTER] key press to terminate the program.
	getchar();

	return status;
}