#include <ntddk.h>

VOID Unload(PDRIVER_OBJECT pDriverObject);
NTSTATUS DriverDispatchWhaterverName(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp);
NTSTATUS DriverEntry(PDRIVER_OBJECT pDriverObject, PUNICODE_STRING pRegistryPath);

#ifdef ALLOC_PRAGMA
#pragma alloc_text(init, DriverEntry)
#pragma alloc_text(page, Unload)
#pragma alloc_text(page, DriverDispatchWhaterverName)
#endif

#define SIOCTL_TYPE 40000
#define IOCTL_Func CTL_CODE(SIOCTL_TYPE, 0x800, METHOD_BUFFERED, FILE_READ_DATA | FILE_WRITE_DATA)

VOID Unload(PDRIVER_OBJECT pDriverObject)
{
    UNREFERENCED_PARAMETER(pDriverObject);
    DbgPrint("Unload routine called.\n");

    // trebuie sa stergi symbolic link-ul creat si device objectul
    IoDeleteDevice(pDriverObject->DeviceObject);

    UNICODE_STRING SymbolicLinkName;
    RtlInitUnicodeString(&SymbolicLinkName, L"\\DosDevices\\MyDevice");


    NTSTATUS DeleteStatus;
    DeleteStatus = IoDeleteSymbolicLink(&SymbolicLinkName);

    if (DeleteStatus != STATUS_SUCCESS)
    {
        return;
    }
}

NTSTATUS DriverDispatchWhaterverName(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    (void)pDeviceObject;
    (void)pIrp;
    DbgPrint("Dispatch called\n");
    return STATUS_SUCCESS;
}

NTSTATUS DriverDispatchIRPMJCreate(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    (void)pDeviceObject;
    (void)pIrp;
    DbgPrint("Create Dispatch called\n");
    return STATUS_SUCCESS;
}

NTSTATUS DriverDispatchIRPMJClose(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    (void)pDeviceObject;
    (void)pIrp;
    DbgPrint("Close Dispatch called\n");
    return STATUS_SUCCESS;
}

NTSTATUS DriverDispatchIRPMJRead(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    (void)pDeviceObject;
    (void)pIrp;
    DbgPrint("Read Dispatch called\n");
    return STATUS_SUCCESS;
}


NTSTATUS DriverDispatchIRPMJWrite(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    PIO_STACK_LOCATION pIoStackIrp = IoGetCurrentIrpStackLocation(pIrp);
    PVOID pWriteDataBuffer = (PVOID)pIrp->AssociatedIrp.SystemBuffer;

    DbgPrint("DriverDispatchIRPMJWrite Called \r\n");
    UNREFERENCED_PARAMETER(pDeviceObject);
    /*
    * Each time the IRP is passed down
    * the driver stack a new stack location is added
    * specifying certain parameters for the IRP to the driver.
    */
    pIoStackIrp = IoGetCurrentIrpStackLocation(pIrp);
    if (pIoStackIrp)
    {
        pWriteDataBuffer = (PVOID)pIrp->AssociatedIrp.SystemBuffer;
        if (pWriteDataBuffer)
        {
            /*
            * Verification needs to be done on the data, bad things
            * happen if we access memory not valid while in the Kernel
            */
            if ((unsigned char)pWriteDataBuffer <= 127)
            {
                DbgPrint(pWriteDataBuffer);
            }
        }
    }
    return NtStatus;
}


NTSTATUS DriverDispatchIRPMJDeviceControl(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pIoStackLocation;
    PCHAR welcome = "Hello from kerneland.";
    PVOID pBuf = pIrp->AssociatedIrp.SystemBuffer;
    DbgPrint("DriverDispatchIRPMJDeviceControl function has been called \n");
    UNREFERENCED_PARAMETER(pDeviceObject);
    pIoStackLocation = IoGetCurrentIrpStackLocation(pIrp);
    switch (pIoStackLocation->Parameters.DeviceIoControl.IoControlCode)
    {
        case IOCTL_Func:
        {
            DbgPrint("IOCTL HELLO.");
            DbgPrint("Message received : %s", pBuf);
            RtlZeroMemory(pBuf,
                pIoStackLocation->Parameters.DeviceIoControl.InputBufferLength
                );
            RtlCopyMemory(pBuf, welcome, strlen(welcome));
            break;
        }
    }
    // Finish the I/O operation by simply completing the packet and returning
    // the same status as in the packet itself.
    pIrp->IoStatus.Status = STATUS_SUCCESS;
    pIrp->IoStatus.Information = strlen(welcome);
    IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}


NTSTATUS
DriverEntry(
PDRIVER_OBJECT pDriverObject,
PUNICODE_STRING pRegistryPath)
{
    DbgPrint("DriverEntry called\n");
    DbgPrint("Registry Path: %wZ\n", pRegistryPath);

    pDriverObject->DriverUnload = Unload;
    UNREFERENCED_PARAMETER(pRegistryPath);

    UNICODE_STRING DriverName, DosDeviceName;
    PDEVICE_OBJECT pDeviceObject;

    RtlInitUnicodeString(&DriverName, L"\\Device\\MYDEVICE");
    RtlInitUnicodeString(&DosDeviceName, L"\\DosDevices\\MyDevice");

    NTSTATUS NtStatus;
    NtStatus = IoCreateDevice(pDriverObject, 0, &DriverName, FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN, FALSE, &pDeviceObject);

    pDeviceObject->Flags |= DO_BUFFERED_IO;
    pDeviceObject->Flags &= (~DO_DEVICE_INITIALIZING);

    if (NtStatus == STATUS_SUCCESS)
    {
        NTSTATUS SymbolicLinkStatus = IoCreateSymbolicLink(&DosDeviceName, &DriverName);

        if (SymbolicLinkStatus != STATUS_SUCCESS)
        {
            return SymbolicLinkStatus;
        }
    }

    for (int i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; i++)
    {
        pDriverObject->MajorFunction[i] = DriverDispatchWhaterverName;
    }

    pDriverObject->MajorFunction[IRP_MJ_CREATE] = DriverDispatchIRPMJCreate;
    pDriverObject->MajorFunction[IRP_MJ_CLOSE] = DriverDispatchIRPMJClose;
    pDriverObject->MajorFunction[IRP_MJ_READ] = DriverDispatchIRPMJRead;
    pDriverObject->MajorFunction[IRP_MJ_WRITE] = DriverDispatchIRPMJWrite;
    pDriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DriverDispatchIRPMJDeviceControl;

    return STATUS_SUCCESS;
}