#include <ntddk.h>
#include <ntddkbd.h>
#include <ntstrsafe.h>
#include "ntifs.h"
#include <windef.h>

#define IMAGE_NUMBEROF_DIRECTORY_ENTRIES    16
#define IMAGE_SIZEOF_SHORT_NAME              8
#define IMAGE_DIRECTORY_ENTRY_IMPORT          1
#define IMAGE_ORDINAL_FLAG32 0x80000000

typedef struct _IMAGE_DATA_DIRECTORY {
    DWORD   VirtualAddress;
    DWORD   Size;
} IMAGE_DATA_DIRECTORY, *PIMAGE_DATA_DIRECTORY;

typedef struct _IMAGE_OPTIONAL_HEADER64 {
    WORD        Magic;
    BYTE        MajorLinkerVersion;
    BYTE        MinorLinkerVersion;
    DWORD       SizeOfCode;
    DWORD       SizeOfInitializedData;
    DWORD       SizeOfUninitializedData;
    DWORD       AddressOfEntryPoint;
    DWORD       BaseOfCode;
    ULONGLONG   ImageBase;
    DWORD       SectionAlignment;
    DWORD       FileAlignment;
    WORD        MajorOperatingSystemVersion;
    WORD        MinorOperatingSystemVersion;
    WORD        MajorImageVersion;
    WORD        MinorImageVersion;
    WORD        MajorSubsystemVersion;
    WORD        MinorSubsystemVersion;
    DWORD       Win32VersionValue;
    DWORD       SizeOfImage;
    DWORD       SizeOfHeaders;
    DWORD       CheckSum;
    WORD        Subsystem;
    WORD        DllCharacteristics;
    ULONGLONG   SizeOfStackReserve;
    ULONGLONG   SizeOfStackCommit;
    ULONGLONG   SizeOfHeapReserve;
    ULONGLONG   SizeOfHeapCommit;
    DWORD       LoaderFlags;
    DWORD       NumberOfRvaAndSizes;
    IMAGE_DATA_DIRECTORY DataDirectory[IMAGE_NUMBEROF_DIRECTORY_ENTRIES];
} IMAGE_OPTIONAL_HEADER64, *PIMAGE_OPTIONAL_HEADER64;

typedef struct _IMAGE_FILE_HEADER {
    WORD    Machine;
    WORD    NumberOfSections;
    DWORD   TimeDateStamp;
    DWORD   PointerToSymbolTable;
    DWORD   NumberOfSymbols;
    WORD    SizeOfOptionalHeader;
    WORD    Characteristics;
} IMAGE_FILE_HEADER, *PIMAGE_FILE_HEADER;

typedef struct _IMAGE_NT_HEADERS64 {
    DWORD Signature;
    IMAGE_FILE_HEADER FileHeader;
    IMAGE_OPTIONAL_HEADER64 OptionalHeader;
} IMAGE_NT_HEADERS64, *PIMAGE_NT_HEADERS64;

typedef IMAGE_NT_HEADERS64                  IMAGE_NT_HEADERS;

#define IMAGE_FIRST_SECTION( ntheader ) ((PIMAGE_SECTION_HEADER)        \
    ((ULONG_PTR)(ntheader)+\
    FIELD_OFFSET(IMAGE_NT_HEADERS, OptionalHeader) + \
    ((ntheader))->FileHeader.SizeOfOptionalHeader   \
    ))

typedef struct _IMAGE_IMPORT_BY_NAME {
    WORD    Hint;
    BYTE    Name[1];
} IMAGE_IMPORT_BY_NAME, *PIMAGE_IMPORT_BY_NAME;

typedef struct _IMAGE_DOS_HEADER {      // DOS .EXE header
    WORD   e_magic;                     // Magic number
    WORD   e_cblp;                      // Bytes on last page of file
    WORD   e_cp;                        // Pages in file
    WORD   e_crlc;                      // Relocations
    WORD   e_cparhdr;                   // Size of header in paragraphs
    WORD   e_minalloc;                  // Minimum extra paragraphs needed
    WORD   e_maxalloc;                  // Maximum extra paragraphs needed
    WORD   e_ss;                        // Initial (relative) SS value
    WORD   e_sp;                        // Initial SP value
    WORD   e_csum;                      // Checksum
    WORD   e_ip;                        // Initial IP value
    WORD   e_cs;                        // Initial (relative) CS value
    WORD   e_lfarlc;                    // File address of relocation table
    WORD   e_ovno;                      // Overlay number
    WORD   e_res[4];                    // Reserved words
    WORD   e_oemid;                     // OEM identifier (for e_oeminfo)
    WORD   e_oeminfo;                   // OEM information; e_oemid specific
    WORD   e_res2[10];                  // Reserved words
    LONG   e_lfanew;                    // File address of new exe header
} IMAGE_DOS_HEADER, *PIMAGE_DOS_HEADER;

typedef struct _IMAGE_SECTION_HEADER {
    BYTE    Name[IMAGE_SIZEOF_SHORT_NAME];
    union {
        DWORD   PhysicalAddress;
        DWORD   VirtualSize;
    } Misc;
    DWORD   VirtualAddress;
    DWORD   SizeOfRawData;
    DWORD   PointerToRawData;
    DWORD   PointerToRelocations;
    DWORD   PointerToLinenumbers;
    WORD    NumberOfRelocations;
    WORD    NumberOfLinenumbers;
    DWORD   Characteristics;
} IMAGE_SECTION_HEADER, *PIMAGE_SECTION_HEADER;

typedef struct _IMAGE_IMPORT_DESCRIPTOR {
    union {
        DWORD   Characteristics;            // 0 for terminating null import descriptor
        DWORD   OriginalFirstThunk;         // RVA to original unbound IAT (PIMAGE_THUNK_DATA)
    };
    DWORD   TimeDateStamp;                  // 0 if not bound,
    // -1 if bound, and real date\time stamp
    //     in IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT (new BIND)
    // O.W. date/time stamp of DLL bound to (Old BIND)

    DWORD   ForwarderChain;                 // -1 if no forwarders
    DWORD   Name;
    DWORD   FirstThunk;                     // RVA to IAT (if bound this IAT has actual addresses)
} IMAGE_IMPORT_DESCRIPTOR;

typedef IMAGE_IMPORT_DESCRIPTOR __unaligned *PIMAGE_IMPORT_DESCRIPTOR;

typedef struct _IMAGE_THUNK_DATA64 {
    union {
        ULONGLONG ForwarderString;  // PBYTE
        ULONGLONG Function;         // PDWORD
        ULONGLONG Ordinal;
        ULONGLONG AddressOfData;    // PIMAGE_IMPORT_BY_NAME
    } u1;
} IMAGE_THUNK_DATA64;
typedef IMAGE_THUNK_DATA64 * PIMAGE_THUNK_DATA64;

typedef PIMAGE_THUNK_DATA64             PIMAGE_THUNK_DATA;

NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath);
VOID UnloadRoutine(IN PDRIVER_OBJECT DriverObject);
NTSTATUS CreateObjectFile(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);
NTSTATUS CloseObjectFile(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);
NTSTATUS DefaultHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);
NTSTATUS DeviceControl(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);

PVOID PsGetProcessSectionBaseAddress(__in PEPROCESS Process);

HANDLE MyOpenProcess(HANDLE ProcessId);

#define THE_BUFFER_LENGTH 100

typedef struct _DEVICE_EXTENSION
{
    unsigned char InBuffer[THE_BUFFER_LENGTH];
}DEVICE_EXTENSION, *PDEVICE_EXTENSION;

//----------------------------------//

NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath)
{
    UNREFERENCED_PARAMETER(RegistryPath);

    DbgPrint("I am in DriverEntry \n");

    UNICODE_STRING DeviceName, DosDeviceName;

    RtlInitUnicodeString(&DeviceName, L"\\Device\\DumpDevice");
    RtlInitUnicodeString(&DosDeviceName, L"\\DosDevices\\DosDumpDevice");

    DriverObject->DriverUnload = UnloadRoutine;

    for (int i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++)
        DriverObject->MajorFunction[i] = DefaultHandler;

    DriverObject->MajorFunction[IRP_MJ_CREATE] = CreateObjectFile;
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = CloseObjectFile;

    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DeviceControl;

    PDEVICE_OBJECT Device;
    NTSTATUS Status = IoCreateDevice(DriverObject, sizeof(DEVICE_EXTENSION), &DeviceName, FILE_DEVICE_UNKNOWN, 0, FALSE, &Device);

    if (!NT_SUCCESS(Status))
        return Status;
    if (!Device)
        return STATUS_UNEXPECTED_IO_ERROR;

    Status = IoCreateSymbolicLink(&DosDeviceName, &DeviceName);

    if (Status != STATUS_SUCCESS)
    {
        DbgPrint("I can't create symbolic link Filter\n");
        return STATUS_FAILED_DRIVER_ENTRY;
    }
    else
        DbgPrint("I am create symbolic link Filter\n");

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

    return STATUS_SUCCESS;
}

VOID UnloadRoutine(IN PDRIVER_OBJECT DriverObject)
{
    DbgPrint("I am in UnloadRoutine \n");

    UNICODE_STRING LineDosDevice;
    RtlInitUnicodeString(&LineDosDevice, L"\\DosDevices\\DosDumpDevice");
    IoDeleteSymbolicLink(&LineDosDevice);

    IoDeleteDevice(DriverObject->DeviceObject);

}

NTSTATUS CreateObjectFile(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    UNREFERENCED_PARAMETER(DeviceObject);

    DbgPrint("I am in CreateObjectFile \n");

    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

NTSTATUS CloseObjectFile(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    UNREFERENCED_PARAMETER(DeviceObject);

    DbgPrint("I am in CloseObjectFile\n");

    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

NTSTATUS DefaultHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    DbgPrint("I am in DefaultHandler\n");

    UNREFERENCED_PARAMETER(DeviceObject);

    Irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return Irp->IoStatus.Status;
}

NTSTATUS DeviceControl(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
    DbgPrint("I am in read DeviceControl.\n");
    PDEVICE_EXTENSION DeviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;

    PIO_STACK_LOCATION IoStackLocation = IoGetCurrentIrpStackLocation(Irp);
    ULONG bufsize = IoStackLocation->Parameters.Read.Length;
    //----------------Наткнулись на буферизированное устройство-----------------------//
    RtlMoveMemory(DeviceExtension->InBuffer, Irp->AssociatedIrp.SystemBuffer, bufsize);
    Irp->IoStatus.Information = bufsize;

    DeviceExtension->InBuffer[bufsize] = '\0';

    DbgPrint("Name of process: %s\n", DeviceExtension->InBuffer);
    HANDLE ProcessId = IoStackLocation->Parameters.DeviceIoControl.IoControlCode;
    // DeviceExtension->PID;

    MyOpenProcess(ProcessId);


    Irp->IoStatus.Status = STATUS_SUCCESS;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return STATUS_SUCCESS;

}

//---------------------Снимаем Дамп--------------------//
HANDLE MyOpenProcess(HANDLE ProcessId)
{
    PEPROCESS Process = NULL;
    NTSTATUS Status;
    HANDLE hProcess = NULL;
    PIMAGE_DOS_HEADER dosHeader;
    PIMAGE_NT_HEADERS64 ntHeader;
    IMAGE_FILE_HEADER header;
    IMAGE_OPTIONAL_HEADER64 opHeader;
    PIMAGE_SECTION_HEADER pSecHeader;
    int i;
    UNREFERENCED_PARAMETER(Process);
    UNREFERENCED_PARAMETER(Status);
    UNREFERENCED_PARAMETER(hProcess);
    UNREFERENCED_PARAMETER(header);

    __debugbreak();

    Status = PsLookupProcessByProcessId(ProcessId, &Process);
    if (Status != STATUS_SUCCESS)
        DbgPrint("I can't in PsLookupProcessByProcessId\n");

    Status = ObOpenObjectByPointer(Process, OBJ_KERNEL_HANDLE, 0, 0x1F0FFF, 0, 0, &hProcess);
    if (Status != STATUS_SUCCESS)
        DbgPrint("I can't in ObOpenObjectByPointer\n");

    //---------------Attach----------------------//
    KeAttachProcess(Process);

    PVOID BaseAddress = PsGetProcessSectionBaseAddress(Process);

    //for (int i = 0; i < 10; i++)
    //PIMAGE_NT_HEADERS64

    //CHAR Buffer[32];
    //ULONG CountOfRead = 0;

    //NtReadVirtualMemory(hProcess, BaseAddress, Buffer, 20, &CountOfRead);
    dosHeader = (PIMAGE_DOS_HEADER)BaseAddress;
    ntHeader = (PIMAGE_NT_HEADERS64)((DWORD)(dosHeader)+(dosHeader->e_lfanew));
    opHeader = ntHeader->OptionalHeader;

    UINT64 pDataDirImportRVA = ntHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
    UINT64 pDataDirImportSize = ntHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size;

    UNREFERENCED_PARAMETER(pDataDirImportSize);

    for (pSecHeader = IMAGE_FIRST_SECTION(ntHeader), i = 0; i < ntHeader->FileHeader.NumberOfSections; i++, pSecHeader++)
    {
        if ((pDataDirImportRVA >= pSecHeader->VirtualAddress) && ((pSecHeader->VirtualAddress + pSecHeader->Misc.VirtualSize) > pDataDirImportRVA))
        {

            UINT64 pIAT = (pDataDirImportRVA - pSecHeader->VirtualAddress) + pSecHeader->PointerToRawData;

            PIMAGE_IMPORT_DESCRIPTOR pImportDescriptor = NULL;
            PIMAGE_THUNK_DATA pImageThunkData = NULL;

            pImportDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)(pIAT + (DWORD)BaseAddress);
            while (pImportDescriptor->FirstThunk != 0)
            {

                //DbgPrint("  Module : %s\n", (pImportDescriptor->Name - pSecHeader->VirtualAddress) + pSecHeader->PointerToRawData + (DWORD)BaseAddress);
                pImageThunkData = (PIMAGE_THUNK_DATA)((DWORD)BaseAddress + pImportDescriptor->FirstThunk - pSecHeader->VirtualAddress + pSecHeader->PointerToRawData);
                while (pImageThunkData->u1.Function)
                {
                    if (!(pImageThunkData->u1.Ordinal & IMAGE_ORDINAL_FLAG32))
                    {
                        PIMAGE_IMPORT_BY_NAME  pStr = (PIMAGE_IMPORT_BY_NAME)((DWORD)BaseAddress + pImageThunkData->u1.AddressOfData - pSecHeader->VirtualAddress + pSecHeader->PointerToRawData);
                        UNREFERENCED_PARAMETER(pStr);
                    }
                    else
                        pImageThunkData++;
                }
                ++pImportDescriptor;
            }
        }

        //-------------------------Пишем в файл----------------------------//

        __debugbreak();
        UNICODE_STRING FileName;
        OBJECT_ATTRIBUTES ObjectAttributes;
        HANDLE HandleFile;
        IO_STATUS_BLOCK StatusBlock;

        RtlInitUnicodeString(&FileName, L"\\DosDevices\\C:\\DumpProcess.txt");

        IO_STATUS_BLOCK StatusBlock2;

        InitializeObjectAttributes(&ObjectAttributes, &FileName, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);

        NTSTATUS Status = ZwCreateFile(&HandleFile, FILE_APPEND_DATA, &ObjectAttributes, &StatusBlock, NULL, FILE_ATTRIBUTE_NORMAL, 0, FILE_OPEN_IF, FILE_SYNCHRONOUS_IO_NONALERT, NULL, 0);

        PBYTE Pointer = BaseAddress;

        Status = ZwWriteFile(HandleFile, NULL, NULL, NULL, &StatusBlock2, Pointer + pSecHeader->VirtualAddress, pSecHeader->Misc.VirtualSize, NULL, NULL);

        ZwClose(HandleFile);
        HandleFile = NULL;
        //-----------------------------------------------------------------//
        DbgPrint("\n%-36s%s", "Section Header name : ", pSecHeader->Name);
        DbgPrint("\n%-36s%#x", "ActualSize of code or data : ", pSecHeader->Misc.VirtualSize);
        DbgPrint("\n%-36s%#x", "Virtual Address(RVA) :", pSecHeader->VirtualAddress);
    }


    DbgPrint("BaseAddress: \n");

    //--------------Deattach---------------------//
    KeDetachProcess();

    ObDereferenceObject(Process);

    return hProcess;
}