#include <ntddk.h>
//#include <ntddkbd.h>
//#include <ntstrsafe.h>
#include <windef.h>
#include <ntifs2.h>
//#include <WinNT.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;

	MyOpenProcess(2200);

	return -1;
}

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_HEADERS ntHeader;
	IMAGE_FILE_HEADER header;
	IMAGE_OPTIONAL_HEADER64 opHeader;
	PIMAGE_SECTION_HEADER pSecHeader;
	int i;
	UNREFERENCED_PARAMETER(Process);
	UNREFERENCED_PARAMETER(Status);
	UNREFERENCED_PARAMETER(hProcess);

	__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_HEADERS)((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;

	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);
					}
					else
					pImageThunkData++;
				}
				++pImportDescriptor;
			}
		}
		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);
	//	printf("\n%-36s%#x", "Size of raw data (rounded to FA) : ", pSecHeader->SizeOfRawData);
	//	printf("\n%-36s%#x", "Pointer to Raw Data : ", pSecHeader->PointerToRawData);
	}


	DbgPrint("BaseAddress: \n");

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

	ObDereferenceObject(Process);

	return hProcess;
}