#include <ntddk.h>
#include <ntifs.h>
#include <ndk/exfuncs.h>
#include <ndk/ketypes.h>
#include <pseh/pseh2.h>
#include <ntstrsafe.h>
#include <debug.h>

NTSTATUS
NTAPI
DriverEntry(IN PDRIVER_OBJECT DriverObject,
            IN PUNICODE_STRING RegPath)
{
    NTSTATUS Status;
    PVOID Buffer, NewBuffer;
    ULONG BufferSize;
    ULONG ReturnLength;
    PSYSTEM_PROCESS_INFORMATION ProcessInfo;
    Buffer = ExAllocatePool(PagedPool,4096);
    if (Buffer){
        BufferSize = 4096;
    }
    else{
        DPRINT1("Exiting, no memory");
        return STATUS_NO_MEMORY;
    }
    while (TRUE)
    {
        ProcessInfo = (PSYSTEM_PROCESS_INFORMATION)(Buffer);
        Status = ZwQuerySystemInformation(SystemProcessInformation,
                                          ProcessInfo,
                                          BufferSize,
                                          &ReturnLength);
        if(Status != STATUS_INFO_LENGTH_MISMATCH && ReturnLength < BufferSize){break;}
        NewBuffer = ExAllocatePool(PagedPool,BufferSize+4096);
        if (!NewBuffer){
            DPRINT1("Exiting, no memory");
            return STATUS_NO_MEMORY;
        }
        ExFreePool(Buffer);
        BufferSize += 4096;
        Buffer = NewBuffer;
    }
    DPRINT1("Student:Vladislav Bermishev\n");
    DPRINT1("Process list:\n");
    while (TRUE){
        if (ProcessInfo->ImageName.Buffer){
            DbgPrint("%i <- %i:%ws\n",(int)ProcessInfo->InheritedFromUniqueProcessId,
                                   (int)ProcessInfo->UniqueProcessId,
                                   ProcessInfo->ImageName.Buffer);
        }
        if (!ProcessInfo->NextEntryOffset) break;
        ProcessInfo = (PVOID)((ULONG_PTR)ProcessInfo + ProcessInfo->NextEntryOffset);
    }
    return STATUS_SUCCESS;
}