#include <ntddk.h>
const int BUFF_SIZE = 512;

NTSTATUS DriverEntry  (IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath);
NTSTATUS MyCreate (IN PDEVICE_OBJECT pDevObj,IN PIRP pIrp);
NTSTATUS MyClose (IN PDEVICE_OBJECT pDevObj, IN PIRP pIrp);
NTSTATUS MyRead (IN PDEVICE_OBJECT pDevObj,IN PIRP pIrp);
NTSTATUS MyWrite (IN PDEVICE_OBJECT pDevObj, IN PIRP pIrp);
NTSTATUS MyIoEvent (IN PDEVICE_OBJECT pDevObj,IN PIRP pIrp);

VOID MyDriverUnload(PDRIVER_OBJECT DriverObject) {
    UNICODE_STRING aux;

    DbgPrint("Unload");

    RtlInitUnicodeString (&aux, L"\\DosDevices\\Djok");
    IoDeleteSymbolicLink (&aux);
    IoDeleteDevice(DriverObject->DeviceObject);
}

NTSTATUS MyCreate (PDEVICE_OBJECT pDevObj, PIRP pIrp) {
    DbgPrint("Create");

    IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

NTSTATUS MyRead (PDEVICE_OBJECT pDevObj,PIRP pIrp) {
    DbgPrint("Read");

    IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

NTSTATUS MyClose (PDEVICE_OBJECT pDevObj, PIRP pIrp) {
    DbgPrint("Close");

    IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

NTSTATUS MyIoEvent (PDEVICE_OBJECT DeviceObject, PIRP pIrp) {
    DbgPrint("IoEvent");

    IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

int parseNext(LPSTR s, int maxSize, int from) {
    int final = from;
    while (final < maxSize && s[final] != '#') ++final;
    s[final] = 0;
    return final;
}

#define __WRITE 0
#define __DELETE 1
#define __READ 2

VOID executeCommand(int op, LPCSTR path, PVOID buff, ULONG len) {
    HANDLE h;
    NTSTATUS status;
    FILE_DISPOSITION_INFORMATION fdi;
    UNICODE_STRING unicodeString;
    ANSI_STRING ansiString;
    OBJECT_ATTRIBUTES objAttr;
    IO_STATUS_BLOCK ioStatusBlock;

    RtlInitUnicodeString(&unicodeString, 0);
    RtlInitAnsiString(&ansiString, path);
    RtlAnsiStringToUnicodeString(&unicodeString, &ansiString, TRUE);

    InitializeObjectAttributes(&objAttr, &unicodeString, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, 0, 0);

    status = ZwCreateFile(
        &h, GENERIC_WRITE, &objAttr, &ioStatusBlock, 0, FILE_ATTRIBUTE_NORMAL,
        0, FILE_OVERWRITE_IF, FILE_SYNCHRONOUS_IO_NONALERT, 0, 0
    );

    if (!NT_SUCCESS(status)) {
        DbgPrint("error CreateFile");
        RtlFreeUnicodeString(&unicodeString);
        return ;
    }

    if (op == __WRITE) {
        status = ZwWriteFile(h, 0, 0, 0, &ioStatusBlock, buff, len, 0, 0);
        if (!NT_SUCCESS(status)) {
            DbgPrint("error Write");
            ZwClose(h);
            RtlFreeUnicodeString(&unicodeString);
            return ;
        }
    }

    if (op == __DELETE) {
        fdi.DeleteFile = TRUE;
        ZwSetInformationFile(h, &ioStatusBlock, &fdi, sizeof(fdi), FileDispositionInformation);
    }

    if (op == __READ) {
        DbgPrint(buff);
    }

    ZwClose(h);
    RtlFreeUnicodeString(&unicodeString);
}

VOID solve(PVOID command) {
    int actionPoz = parseNext((LPSTR)command, BUFF_SIZE, 0);
    int pathPoz = parseNext((LPSTR)command, BUFF_SIZE, actionPoz + 1);
    int paramsPoz = parseNext((LPSTR)command, BUFF_SIZE, pathPoz + 1);

    LPCSTR path = ((LPCSTR)command) + actionPoz + 1;
    PVOID params = (PVOID)(((LPCSTR )command) + pathPoz + 1);
    int paramSize = paramsPoz - pathPoz - 1;
    DbgPrint(path);
    DbgPrint((LPCSTR)params);
    DbgPrint((LPCSTR)command);

    if(((LPCSTR)command)[0] == 'C') executeCommand(__WRITE, path, (PVOID)"", 0);
    if(((LPCSTR)command)[0] == 'W') executeCommand(__WRITE, path, params, paramSize);
    if(((LPCSTR)command)[0] == 'D') executeCommand(__DELETE, path, (PVOID)"", 0);
    if(((LPCSTR)command)[0] == 'R') executeCommand(__READ, path, (PVOID)params, paramSize);
}

NTSTATUS MyWrite (PDEVICE_OBJECT pDevObj, PIRP pIrp) {
    DbgPrint("Write");

    if(pIrp->AssociatedIrp.SystemBuffer!=NULL) {
        DbgPrint(pIrp->AssociatedIrp.SystemBuffer);
        solve(pIrp->AssociatedIrp.SystemBuffer);
    }

    IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

NTSTATUS DriverEntry(PDRIVER_OBJECT driverObj, PUNICODE_STRING regPath) {
    NTSTATUS status;
    UNICODE_STRING usDriverName;
    UNICODE_STRING usDosDeviceName;
    PDEVICE_OBJECT pDeviceObj = 0;

    RtlInitUnicodeString(&usDriverName, L"\\Device\\Djok");
    RtlInitUnicodeString(&usDosDeviceName, L"\\DosDevices\\Djok");

    status = IoCreateDevice(driverObj, 0, &usDriverName, FILE_DEVICE_UNKNOWN, 0, 0, &pDeviceObj);
    if(!NT_SUCCESS(status))
        return status;

    pDeviceObj->Flags |= DO_BUFFERED_IO;
    pDeviceObj->Flags &= ~DO_DEVICE_INITIALIZING;

    status = IoCreateSymbolicLink(&usDosDeviceName, &usDriverName);
    if(!NT_SUCCESS(status)) {
        IoDeleteDevice(driverObj->DeviceObject);
        return status;
    }

    driverObj->MajorFunction[IRP_MJ_CREATE] = MyCreate;
    driverObj->MajorFunction[IRP_MJ_CLOSE] = MyClose;
    driverObj->MajorFunction[IRP_MJ_READ] = MyRead;
    driverObj->MajorFunction[IRP_MJ_WRITE] = MyWrite;
    driverObj->MajorFunction[IRP_MJ_DEVICE_CONTROL] = MyIoEvent;
    driverObj->DriverUnload = MyDriverUnload;

    DbgPrint("Hello");
    return STATUS_SUCCESS;
}