Integration test improved.

1 files changed, 122 insertions, 94 deletions

diff --git a/IntegrationTest/IntegrationTest.cpp b/IntegrationTest/IntegrationTest.cpp
index 6a53d94..b4f852e 100644
--- a/IntegrationTest/IntegrationTest.cpp
+++ b/IntegrationTest/IntegrationTest.cpp
@@ -1,6 +1,3 @@
-// IntegrationTest.cpp : This file contains the 'main' function. Program execution begins and ends there.
-//
-
 #include "pch.h"
 #include "KInterface.h"
 
@@ -12,7 +9,124 @@
 #define KM_ASSERT_EQUAL(equal, condition, message) \
 	do { PRINT_CHECK_MSG(message); if ((condition) != (equal)) { \
 		PRINT_FAIL_MSG(); goto error; } else { PRINT_OK_MSG(); } \
-	} while (0);
+	} while (0)
+#define KM_TEST_SUITE(condition, message) \
+	do { \
+		PRINT_CHECK_MSG("--- TestSuite " ## message ## " ---\n"); \
+		if ((condition) != true) { goto error; } \
+	} while(0)
+
+static bool test_Processes(KInterface& ki)
+{
+	SIZE_T required_processes_found = 0;
+	std::vector<PROCESS_DATA> processes;
+	KM_ASSERT_EQUAL(true, ki.Processes(processes), "Kernel Interface Processes");
+	KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
+
+	for (auto& process : processes)
+	{
+		//std::cout << "Process Name: " << process.ImageName << std::endl;
+		if (strcmp(process.ImageName, "IntegrationTest-kmem.exe") == 0 && strlen(process.ImageName) == strlen("IntegrationTest-kmem.exe"))
+		{
+			required_processes_found++;
+		}
+		if (strcmp(process.ImageName, "System") == 0 && strlen(process.ImageName) == strlen("System"))
+		{
+			required_processes_found++;
+		}
+		if (strcmp(process.ImageName, "Registry") == 0 && strlen(process.ImageName) == strlen("Registry"))
+		{
+			required_processes_found++;
+		}
+		if (strcmp(process.ImageName, "wininit.exe") == 0 && strlen(process.ImageName) == strlen("wininit.exe"))
+		{
+			required_processes_found++;
+		}
+		if (strcmp(process.ImageName, "winlogon.exe") == 0 && strlen(process.ImageName) == strlen("winlogon.exe"))
+		{
+			required_processes_found++;
+		}
+		if (strcmp(process.ImageName, "lsass.exe") == 0 && strlen(process.ImageName) == strlen("lsass.exe"))
+		{
+			required_processes_found++;
+		}
+	}
+	KM_ASSERT_EQUAL(6, required_processes_found, "Kernel Interface Modules (6 required found)");
+
+	return true;
+error:
+	return false;
+}
+
+static bool test_Modules(KInterface& ki, HANDLE pid)
+{
+	SIZE_T required_modules_found = 0;
+	std::vector<MODULE_DATA> modules;
+	KM_ASSERT_EQUAL(true, ki.Modules(pid, modules), "Kernel Interface Modules");
+	KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
+	for (auto& module : modules) {
+		//std::cout << "DLL Name: " << module.BaseDllName << std::endl;
+		if (strcmp(module.BaseDllName, "IntegrationTest-kmem.exe") == 0 && strlen(module.BaseDllName) == strlen("IntegrationTest-kmem.exe"))
+		{
+			required_modules_found++;
+		}
+		if (strcmp(module.BaseDllName, "ntdll.dll") == 0 && strlen(module.BaseDllName) == strlen("ntdll.dll"))
+		{
+			required_modules_found++;
+		}
+		if (strcmp(module.BaseDllName, "KERNEL32.DLL") == 0 && strlen(module.BaseDllName) == strlen("KERNEL32.DLL"))
+		{
+			required_modules_found++;
+		}
+	}
+	KM_ASSERT_EQUAL(3, required_modules_found, "Kernel Interface Modules (3 required found)");
+
+	return true;
+error:
+	return false;
+}
+
+static bool test_Pages(KInterface& ki, HANDLE pid)
+{
+	SIZE_T found_shmaddr = 0;
+	std::vector<MEMORY_BASIC_INFORMATION> pages;
+	KM_ASSERT_EQUAL(true, ki.Pages(pid, pages), "Kernel Interface Pages");
+	KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
+	for (auto& page : pages) {
+		if (page.BaseAddress == (PVOID)SHMEM_ADDR && page.RegionSize == SHMEM_SIZE) {
+			found_shmaddr++;
+		}
+	}
+	KM_ASSERT_EQUAL(1, found_shmaddr, "Kernel Interface Pages (1 required found)");
+
+	return true;
+error:
+	return false;
+}
+
+static bool test_VirtualMemory(KInterface& ki, HANDLE pid)
+{
+	PVOID addr = (PVOID)SHMEM_ADDR;
+	SIZE_T size = 0x100;
+
+	KM_ASSERT_EQUAL(true,
+		ki.VAlloc(pid, &addr, &size, PAGE_READWRITE), "Kernel Interface VirtualAlloc SHMEM_ADDR");
+	KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
+
+	addr = NULL;
+	size = 0x100;
+
+	KM_ASSERT_EQUAL(true,
+		ki.VAlloc(pid, &addr, &size, PAGE_READWRITE), "Kernel Interface VirtualAlloc");
+	KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
+	KM_ASSERT_EQUAL(true,
+		ki.VFree(pid, addr, size), "Kernel Interface VirtualFree");
+	KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
+
+	return true;
+error:
+	return false;
+}
 
 int main()
 {
@@ -30,96 +144,10 @@ int main()
 		KM_ASSERT_EQUAL(true, ki.Ping(), "Kernel Interface PING - PONG #1");
 		KM_ASSERT_EQUAL(true, ki.Ping(), "Kernel Interface PING - PONG #2");
 		KM_ASSERT_EQUAL(true, ki.Ping(), "Kernel Interface PING - PONG #3");
-
-		{
-			SIZE_T required_processes_found = 0;
-			std::vector<PROCESS_DATA> processes;
-			KM_ASSERT_EQUAL(true, ki.Processes(processes), "Kernel Interface Processes");
-			KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
-
-			for (auto& process : processes)
-			{
-				//std::cout << "Process Name: " << process.ImageName << std::endl;
-				if (strcmp(process.ImageName, "IntegrationTest-kmem.exe") == 0 && strlen(process.ImageName) == strlen("IntegrationTest-kmem.exe"))
-				{
-					required_processes_found++;
-				}
-				if (strcmp(process.ImageName, "System") == 0 && strlen(process.ImageName) == strlen("System"))
-				{
-					required_processes_found++;
-				}
-				if (strcmp(process.ImageName, "Registry") == 0 && strlen(process.ImageName) == strlen("Registry"))
-				{
-					required_processes_found++;
-				}
-				if (strcmp(process.ImageName, "wininit.exe") == 0 && strlen(process.ImageName) == strlen("wininit.exe"))
-				{
-					required_processes_found++;
-				}
-				if (strcmp(process.ImageName, "winlogon.exe") == 0 && strlen(process.ImageName) == strlen("winlogon.exe"))
-				{
-					required_processes_found++;
-				}
-				if (strcmp(process.ImageName, "lsass.exe") == 0 && strlen(process.ImageName) == strlen("lsass.exe"))
-				{
-					required_processes_found++;
-				}
-			}
-			KM_ASSERT_EQUAL(6, required_processes_found, "Kernel Interface Modules (6 required found)");
-		}
-
-		{
-			SIZE_T required_modules_found = 0;
-			std::vector<MODULE_DATA> modules;
-			KM_ASSERT_EQUAL(true, ki.Modules(this_pid, modules), "Kernel Interface Modules");
-			KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
-			for (auto& module : modules) {
-				//std::cout << "DLL Name: " << module.BaseDllName << std::endl;
-				if (strcmp(module.BaseDllName, "IntegrationTest-kmem.exe") == 0 && strlen(module.BaseDllName) == strlen("IntegrationTest-kmem.exe"))
-				{
-					required_modules_found++;
-				}
-				if (strcmp(module.BaseDllName, "ntdll.dll") == 0 && strlen(module.BaseDllName) == strlen("ntdll.dll"))
-				{
-					required_modules_found++;
-				}
-				if (strcmp(module.BaseDllName, "KERNEL32.DLL") == 0 && strlen(module.BaseDllName) == strlen("KERNEL32.DLL"))
-				{
-					required_modules_found++;
-				}
-			}
-			KM_ASSERT_EQUAL(3, required_modules_found, "Kernel Interface Modules (3 required found)");
-		}
-
-		{
-			SIZE_T found_shmaddr = 0;
-			std::vector<MEMORY_BASIC_INFORMATION> pages;
-			KM_ASSERT_EQUAL(true, ki.Pages(this_pid, pages), "Kernel Interface Pages");
-			KM_ASSERT_EQUAL(0, ki.getLastNtStatus(), "Last NtStatus");
-			for (auto& page : pages) {
-				if (page.BaseAddress == (PVOID)SHMEM_ADDR && page.RegionSize == SHMEM_SIZE) {
-					found_shmaddr++;
-				}
-			}
-			KM_ASSERT_EQUAL(1, found_shmaddr, "Kernel Interface Pages (1 required found)");
-		}
-
-		{
-			PVOID addr = (PVOID)SHMEM_ADDR;
-			SIZE_T size = 0x100;
-
-			KM_ASSERT_EQUAL(true,
-				ki.VAlloc(this_pid, &addr, &size, PAGE_READWRITE), "Kernel Interface VirtualAlloc SHMEM_ADDR");
-
-			addr = NULL;
-			size = 0x100;
-
-			KM_ASSERT_EQUAL(true,
-				ki.VAlloc(this_pid, &addr, &size, PAGE_READWRITE), "Kernel Interface VirtualAlloc");
-			KM_ASSERT_EQUAL(true,
-				ki.VFree(this_pid, addr, size), "Kernel Interface VirtualFree");
-		}
-
+		KM_TEST_SUITE(test_Processes(ki), "Processes");
+		KM_TEST_SUITE(test_Modules(ki, this_pid), "Modules");
+		KM_TEST_SUITE(test_Pages(ki, this_pid), "Pages");
+		KM_TEST_SUITE(test_VirtualMemory(ki, this_pid), "VirtualMemory");
 		KM_ASSERT_EQUAL(true, ki.Ping(), "Kernel Interface PING - PONG #4");
 		KM_ASSERT_EQUAL(true, ki.Exit(), "Kernel Interface Driver Shutdown");
 	}