#include <ntifs.h>
#include <ntddk.h>
#include "ThreadData.h"

#define DRIVER_TAG 'dcba'

UNICODE_STRING g_RegistryPath;

NTSTATUS CreateClose(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);
NTSTATUS DeviceControl(IN PDEVICE_OBJECT, IN PIRP Irp);
void SampleUnload(IN PDRIVER_OBJECT DriverObject);

extern "C"
NTSTATUS
DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath) {
	DriverObject->DriverUnload = SampleUnload;

	g_RegistryPath.Buffer = (WCHAR*)ExAllocatePool2(
		PagedPool,
		RegistryPath->Length,
		DRIVER_TAG);

	if (g_RegistryPath.Buffer == nullptr) {
		KdPrint(("Failed to allocate memory: %u\n", RegistryPath->Length));
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	g_RegistryPath.MaximumLength = RegistryPath->Length;

	RtlCopyUnicodeString(&g_RegistryPath, RegistryPath);

	KdPrint(("Sample driver initialized successfully\n"));

	KdPrint(("Copied registry path: %wZ\n", &g_RegistryPath));

	DriverObject->MajorFunction[IRP_MJ_CREATE] = CreateClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = CreateClose;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DeviceControl;

	UNICODE_STRING devName = RTL_CONSTANT_STRING(L"\\Device\\Sample");

	PDEVICE_OBJECT DeviceObject;

	NTSTATUS status = IoCreateDevice(
		DriverObject,
		0,
		&devName,
		FILE_DEVICE_UNKNOWN,
		0,
		FALSE,
		&DeviceObject);

	if (!NT_SUCCESS(status)) {
		KdPrint(("Failed to create device object (0x%08X)\n", status));
		return status;
	}

	UNICODE_STRING symLink = RTL_CONSTANT_STRING(L"\\??\\Sample");
	status = IoCreateSymbolicLink(&symLink, &devName);
	if (!NT_SUCCESS(status)) {
		KdPrint(("Failed to create symbolic link (0x%08X)\n", status));
		IoDeleteDevice(DeviceObject);
		return status;
	}

	return STATUS_SUCCESS;
}

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


NTSTATUS DeviceControl(IN PDEVICE_OBJECT, IN PIRP Irp) {
	auto stack = IoGetCurrentIrpStackLocation(Irp);
	auto status = STATUS_SUCCESS;

	switch (stack->Parameters.DeviceIoControl.IoControlCode) {
		case IOCTL_SET_PRIORITY: {
			auto len = stack->Parameters.DeviceIoControl.InputBufferLength;
			if (len < sizeof(ThreadData)) {
				status = STATUS_BUFFER_TOO_SMALL;
				break;
			}
			auto data = (ThreadData*)stack->Parameters.DeviceIoControl.Type3InputBuffer;
			if (data == nullptr) {
				status = STATUS_INVALID_PARAMETER;
				break;
			}
			if (data->Priority < 1 || data->Priority > 31) {
				status = STATUS_INVALID_PARAMETER;
				break;
			}

			PETHREAD Thread;

			status = PsLookupThreadByThreadId(ULongToHandle(data->ThreadId), &Thread);
			if (!NT_SUCCESS(status))
				break;

			KeSetPriorityThread((PKTHREAD)Thread, data->Priority);

			ObDereferenceObject(Thread);

			KdPrint(("Thread Priority change for %d to %d succeeded!\n",
				data->ThreadId, data->Priority));

			break;
		}
		default:
			status = STATUS_INVALID_DEVICE_REQUEST;
			break;
	}

	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = 0;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return status;
}

void SampleUnload(IN PDRIVER_OBJECT DriverObject) {
	ExFreePool(g_RegistryPath.Buffer);

	UNICODE_STRING symLink = RTL_CONSTANT_STRING(L"\\??\\Sample");
	IoDeleteSymbolicLink(&symLink);

	IoDeleteDevice(DriverObject->DeviceObject);

	KdPrint(("Sample driver Unload called\n"));
}