注册表回调与注册表保护

#include <ntifs.h>
#include<ntddk.h>
#include <stdio.h>
#include<stdlib.h>
#include< Ntstrsafe.h>


NTKERNELAPI
PCHAR
PsGetProcessImageFileName(
	IN PEPROCESS Process
);

NTSTATUS GetRegistryObjectCompleteName(
	OUT PUNICODE_STRING* RegistryPath,
	IN PVOID RegistryObject
)
{
	NTSTATUS		Status;
	ULONG			NeededLength;
	POBJECT_NAME_INFORMATION	ObjectName = NULL;


	// 确保对象内存可访问
	if ((!MmIsAddressValid(RegistryObject)) ||
		(RegistryObject == NULL))
	{
		return STATUS_UNSUCCESSFUL;
	}

	// 先获取注册表对象名所占用的内存大小
	Status = ObQueryNameString(RegistryObject, NULL, 0, &NeededLength);

	if (Status == STATUS_INFO_LENGTH_MISMATCH)
	{
		// 申请内存 多2个字节方便比对
		ObjectName = ExAllocatePoolWithTag(PagedPool, (ULONG64)NeededLength + 2, 'reg');

		if (!ObjectName)
		{
			return STATUS_UNSUCCESSFUL;
		}
		RtlZeroMemory(ObjectName, (ULONG64)NeededLength + 2);

		// 获取名字
		Status = ObQueryNameString(RegistryObject, (POBJECT_NAME_INFORMATION)ObjectName, NeededLength, &NeededLength);
		*RegistryPath = &ObjectName->Name;
	}
	return Status;
}

NTSTATUS RegNtSetValueKeyCallBak(IN PVOID Argument2)
{
	NTSTATUS Status = 0;
	PREG_SET_VALUE_KEY_INFORMATION 	SetKey = (PREG_SET_VALUE_KEY_INFORMATION)Argument2;
	PUNICODE_STRING					RegistryFullPath = NULL;
	WCHAR* Path = NULL;

	do
	{
		if (PsGetCurrentProcessId() == (HANDLE)4)
		{
			return 0;
		}
		if (strstr(PsGetProcessImageFileName(PsGetCurrentProcess()), "mute") == NULL)
		{
			return 0;
		}

		if (!NT_SUCCESS(GetRegistryObjectCompleteName(&RegistryFullPath, SetKey->Object)))
		{
			break;
		}

		//加上0结尾
		ULONG64 Length = (ULONG64)RegistryFullPath->Length + 2;
		Path = ExAllocatePoolWithTag(PagedPool, Length, 'reg');
		if (!Path)
		{
			break;
		}
		RtlZeroMemory(Path, Length);
		RtlCopyMemory(Path, RegistryFullPath->Buffer, RegistryFullPath->Length);

		switch (SetKey->Type)
		{
		case REG_DWORD:
		{
			ULONG Data = 0;
			Data = *(ULONG*)SetKey->Data;
			DbgPrint("REG_DWORD 当前进程名:%s  设置注册表项:%wZ 键值:%wZ  设置的值 %d\n", PsGetProcessImageFileName(PsGetCurrentProcess()), RegistryFullPath, SetKey->ValueName, Data);
			break;	
		}
		case REG_SZ:
		{
			DbgPrint("REG_SZ 当前进程名:%s  设置注册表项:%wZ 键值:%wZ  设置的值 %S\n", PsGetProcessImageFileName(PsGetCurrentProcess()), RegistryFullPath, SetKey->ValueName, (WCHAR*)SetKey->Data);
			break;
		}

		default:
			DbgPrint("未处理类型 %d\n", SetKey->Type);
			break;
		}
		
	} while (FALSE);

	if (Path)
	{
		ExFreePool(Path);
		Path = NULL;
	}
	if (RegistryFullPath)
	{
		ExFreePool(RegistryFullPath);
		RegistryFullPath = NULL;
	}
	return Status;
}

// 注册表回调函数的处理 即注册表保护的实现
NTSTATUS
RegistryCallback(
	IN PVOID CallbackContext,
	IN PVOID Argument1,
	IN PVOID Argument2
)
{
	REG_NOTIFY_CLASS	Type;
	NTSTATUS			Status = STATUS_SUCCESS;

	UNREFERENCED_PARAMETER(CallbackContext);

	// 获取操作类型
	Type = (REG_NOTIFY_CLASS)Argument1;

	switch (Type)
	{
	case RegNtSetValueKey:
	{
		RegNtSetValueKeyCallBak(Argument2);
	}
	break;
	}
	return Status;
}
LARGE_INTEGER Cookie;
VOID DriverUnload(IN PDRIVER_OBJECT pDriverObject)
{

	CmUnRegisterCallback(Cookie);
	KdPrint(("驱动卸载成功\n"));
}

NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject, IN PUNICODE_STRING RegistryPath)
{
	NTSTATUS status = STATUS_SUCCESS;
	pDriverObject->DriverUnload = DriverUnload;
	DbgPrint("驱动加载成功");

	UNICODE_STRING Attitude = RTL_CONSTANT_STRING(L"310000");
	status = CmRegisterCallbackEx(RegistryCallback, &Attitude, pDriverObject, NULL, &Cookie, NULL);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("注册表回调函数注册失败\n"));
		return status;
	}
	return status;
}