Merge remote-tracking branch 'impl/VS-2017'

4 files changed, 32 insertions, 14 deletions

diff --git a/KMemDriver/Crypto.c b/KMemDriver/Crypto.c
index 1999d5f..853623a 100644
--- a/KMemDriver/Crypto.c
+++ b/KMemDriver/Crypto.c
@@ -1,15 +1,30 @@
 #include "Crypto.h"
 
+#include <stdarg.h>
+
 struct crypt_data {
 	UINT64 key;
 	UINT8 crypted;
-	UINT32 marker;
+	UINT8 used;
 };
 
-#define MAX_CRYPTED_FUNCTIONS 64
-static struct crypt_data data[MAX_CRYPTED_FUNCTIONS];
+static struct crypt_data * data = NULL;
 static size_t data_used = 0;
 
-void crypt_fn(void)
+void CryptoInit(PVOID fn, ...)
+{
+	va_list ap;
+
+	va_start(ap, fn);
+	PVOID f;
+	while ((f = va_arg(ap, PVOID)) != NULL)
+	{
+
+	}
+	va_end(ap);
+}
+
+void CryptoDo(PVOID fn)
 {
+	UNREFERENCED_PARAMETER(fn);
 }
\ No newline at end of file
diff --git a/KMemDriver/Crypto.h b/KMemDriver/Crypto.h
index 09413aa..2d81b3d 100644
--- a/KMemDriver/Crypto.h
+++ b/KMemDriver/Crypto.h
@@ -2,11 +2,16 @@
 
 #include <ntddk.h>
 
-void crypt_fn(void);
+extern PVOID getNextRIP(void);
+
+void CryptoInit(PVOID fn, ...);
+void CryptoDo(PVOID fn);
+
+#define CRYPTO_FNPTR(fn) ((PVOID)fn)
 
 #define CRYPT_PROLOGUE() \
 	do { \
-		crypt_fn(); \
+		CryptoDo(getNextRIP()); \
 		volatile UINT64 index_and_marker = { 0x11111111C0DEC0DE }; \
 		UNREFERENCED_PARAMETER(index_and_marker); \
 	} while (0)
@@ -14,5 +19,5 @@ void crypt_fn(void);
 	do { \
 		volatile UINT32 marker = 0xDEADDEAD;\
 		UNREFERENCED_PARAMETER(marker); \
-		crypt_fn(); \
+		CryptoDo(getNextRIP()); \
 	} while (0)
\ No newline at end of file
diff --git a/KMemDriver/KMemDriver.c b/KMemDriver/KMemDriver.c
index 921ab0f..db865d1 100644
--- a/KMemDriver/KMemDriver.c
+++ b/KMemDriver/KMemDriver.c
@@ -27,8 +27,6 @@
 
 #define WAIT_OBJECT_0 ((STATUS_WAIT_0 ) + 0 )
 
-extern PVOID getCurrentRIP(void);
-
 DRIVER_INITIALIZE DriverEntry;
 #pragma alloc_text(INIT, DriverEntry)
 void OnImageLoad(
@@ -165,6 +163,7 @@ NTSTATUS DriverEntry(
 	_In_  PUNICODE_STRING RegistryPath
 )
 {
+	CryptoInit(CRYPTO_FNPTR(DriverEntry), NULL);
 	CRYPT_PROLOGUE();
 	NTSTATUS status;
 	HANDLE hThread = NULL;
@@ -175,6 +174,7 @@ NTSTATUS DriverEntry(
 	UNREFERENCED_PARAMETER(RegistryPath);
 
 	KDBG("Driver Loaded\n");
+
 	if (!DriverObject && RegistryPath) {
 		/* assume that we are manual mapped by PastDSE */
 		mmapedBase = RegistryPath;
@@ -198,8 +198,6 @@ NTSTATUS WaitForControlProcess(OUT PEPROCESS *ppEProcess)
 	if (!ppEProcess)
 		return STATUS_INVALID_ADDRESS;
 
-	KDBG("CurrentRIP: %p\n", getCurrentRIP());
-
 	imageBase = NULL;
 	ctrlPID = NULL;
 
diff --git a/KMemDriver/Utils.asm b/KMemDriver/Utils.asm
index d95a64e..b7c344e 100644
--- a/KMemDriver/Utils.asm
+++ b/KMemDriver/Utils.asm
@@ -1,11 +1,11 @@
-PUBLIC getCurrentRIP
+PUBLIC getNextRIP
 
 .code _text
 
-getCurrentRIP PROC PUBLIC
+getNextRIP PROC PUBLIC
 pop rax
 push rax
 ret
-getCurrentRIP ENDP
+getNextRIP ENDP
 
 END
\ No newline at end of file