ProcProt.c

#include <ntddk.h>
#include <windef.h>

#define DEVNAME L"\\Device\\ProcProt"
#define LINKNAME L"\\??\\ProcProt"

PVOID regHandle;

#define IOCTL_PROCESS_PROTECT_BY_PID CTL_CODE(FILE_DEVICE_UNKNOWN , 1, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_PROCESS_UNPROTECT_BY_PID CTL_CODE(FILE_DEVICE_UNKNOWN , 2, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_PROCESS_PROTECT_CLEAR CTL_CODE(FILE_DEVICE_UNKNOWN , 3, METHOD_BUFFERED, FILE_ANY_ACCESS)

void PrUnload(PDRIVER_OBJECT DriverObject) {
	NTSTATUS status;
	UNICODE_STRING sLinkName;
	PDEVICE_OBJECT DevObj, t;

	DbgPrint("PrUnload called \n");

	ObUnRegisterCallbacks(regHandle);

	RtlInitUnicodeString(&sLinkName, LINKNAME);
	status = IoDeleteSymbolicLink(&sLinkName);
	if (status != STATUS_SUCCESS) {
		DbgPrint("IoDeleteSymbolicLink() failed ??!?\n");
	}

	DevObj = DriverObject->DeviceObject;

	while (DevObj) {
		t = DevObj->NextDevice;
		IoDeleteDevice(DevObj);
		DevObj = t;
	}

	DbgPrint("PrUnload is done, module unloaded \n");
	return;
}

#define MAX_PID_ARRAY 1024

HANDLE ProtPids[MAX_PID_ARRAY];
FAST_MUTEX pidMutex;

BOOL isPidAlreadySet(HANDLE pid) {
	unsigned int i;
	for (i = 0; i < MAX_PID_ARRAY; i++) {
		if (ProtPids[i] == pid) return TRUE;
	}

	return FALSE;
}

BOOL IsPidProtected(HANDLE pid) {
//	unsigned int i;
	BOOL r = FALSE;;
	ExAcquireFastMutex(&pidMutex);
	r = isPidAlreadySet(pid);
	ExReleaseFastMutex(&pidMutex);
	return r;
}

void addPid(HANDLE pid) {
	unsigned int i;
	ExAcquireFastMutex(&pidMutex);

	if (isPidAlreadySet(pid)) goto END;

	for (i = 0; i < MAX_PID_ARRAY; i++) {
		if (ProtPids[i] == NULL) {
			ProtPids[i] = pid;
			goto END;
		}
	}

END:
	ExReleaseFastMutex(&pidMutex);
	return;
}

void removePid(HANDLE pid) {
	unsigned int i;
	ExAcquireFastMutex(&pidMutex);

	for (i = 0; i < MAX_PID_ARRAY; i++) {
		if (ProtPids[i] == pid) {
			ProtPids[i] = NULL;
			goto END;
		}
	}

END:
	ExReleaseFastMutex(&pidMutex);
	return;
}

void clearPids() {
	ExAcquireFastMutex(&pidMutex);
	memset(ProtPids, 0, sizeof(ProtPids));
	ExReleaseFastMutex(&pidMutex);
	return;
}

NTSTATUS CreateCloseDispatch(
	PDEVICE_OBJECT DeviceObject,
	PIRP Irp
) {
	UNREFERENCED_PARAMETER(DeviceObject);
	Irp->IoStatus.Information = 0;
	Irp->IoStatus.Status = STATUS_SUCCESS;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS IoctlDispatch(
	PDEVICE_OBJECT DeviceObject,
	PIRP Irp
) {
	PIO_STACK_LOCATION IrpSp = IoGetCurrentIrpStackLocation(Irp);
	NTSTATUS status = STATUS_SUCCESS;
	DWORD info = 0;
	HANDLE pid;

	UNREFERENCED_PARAMETER(DeviceObject);

	switch (IrpSp->Parameters.DeviceIoControl.IoControlCode) {
	case IOCTL_PROCESS_PROTECT_BY_PID:
		if (IrpSp->Parameters.DeviceIoControl.InputBufferLength != sizeof(HANDLE)) {
			status = STATUS_INVALID_BUFFER_SIZE;
			goto END;

		}

		pid = *(PHANDLE) Irp->AssociatedIrp.SystemBuffer;
		addPid(pid);
		info = 0;
		break;

	case IOCTL_PROCESS_UNPROTECT_BY_PID: 
		if (IrpSp->Parameters.DeviceIoControl.InputBufferLength != sizeof(HANDLE)) {
			status = STATUS_INVALID_BUFFER_SIZE;
			goto END;

		}

		pid = *(PHANDLE)Irp->AssociatedIrp.SystemBuffer;
		removePid(pid);
		info = 0;
		break;

	case IOCTL_PROCESS_PROTECT_CLEAR:
		clearPids();
		info = 0;
		break;
	}

END:
	Irp->IoStatus.Information = info;
	Irp->IoStatus.Status = status;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;

}

#define PROCESS_TERMINATE 0x0001

OB_PREOP_CALLBACK_STATUS OnPreOpenProcess(
	PVOID RegistrationContext,
	POB_PRE_OPERATION_INFORMATION OperationInformation
)
{
	UNREFERENCED_PARAMETER(RegistrationContext);

	// both create and duplicate first arg is DesiredAccess, so we can optimize .... 
	if ((OperationInformation->Parameters->CreateHandleInformation.DesiredAccess & PROCESS_TERMINATE) == 0) goto END;

	HANDLE pid = PsGetProcessId((PEPROCESS)OperationInformation->Object);
	if (!IsPidProtected(pid)) goto END;

	OperationInformation->Parameters->CreateHandleInformation.DesiredAccess &= ~PROCESS_TERMINATE;

END:
	return OB_PREOP_SUCCESS;
}

OB_OPERATION_REGISTRATION operations[] = {
	{
		/*PsProcessType*/ NULL,
		OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE,
		OnPreOpenProcess,
		NULL
	}
};

OB_CALLBACK_REGISTRATION reg = {
	OB_FLT_REGISTRATION_VERSION, 1
};


NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath) {
	PDEVICE_OBJECT DevObject;
	UNICODE_STRING sDevName;
	UNICODE_STRING sLinkName;
	NTSTATUS status;

	DbgPrint("Simple DriverEntry called: %wZ\n", RegistryPath);

	DriverObject->DriverUnload = PrUnload;

	RtlInitUnicodeString(&sDevName, DEVNAME);
	RtlInitUnicodeString(&sLinkName, LINKNAME);

	status = IoCreateDevice(DriverObject, 0, &sDevName, FILE_DEVICE_UNKNOWN, 0, FALSE, &DevObject);
	if (status != STATUS_SUCCESS) {
		DbgPrint("IoCreateDevice() failed\n");
		return STATUS_UNSUCCESSFUL;
	}

	status = IoCreateSymbolicLink(&sLinkName, &sDevName);
	if (status != STATUS_SUCCESS) {
		DbgPrint("IoCreateSymbolicLink() failed\n");
		IoDeleteDevice(DevObject);
		return STATUS_UNSUCCESSFUL;
	}

	DriverObject->MajorFunction[IRP_MJ_CREATE] = CreateCloseDispatch;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = CreateCloseDispatch;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = IoctlDispatch;


	regHandle = NULL;

	operations[0].ObjectType = PsProcessType;

	reg.OperationRegistration = operations;
	RtlInitUnicodeString(&reg.Altitude, L"12345.6171");

	status = ObRegisterCallbacks(&reg, &regHandle);
	if (status != STATUS_SUCCESS) {
		DbgPrint("ObRegisterCallbacks() failed\n");
		IoDeleteDevice(DevObject);
		IoDeleteSymbolicLink(&sLinkName);
		return STATUS_UNSUCCESSFUL;
	}

	memset(&ProtPids, 0, sizeof(ProtPids));
	ExInitializeFastMutex(&pidMutex);

	DevObject->Flags &= ~DO_DEVICE_INITIALIZING;

	DbgPrint("Driver is successfully loaded!\n");
	return STATUS_SUCCESS;
}