/*++

block Name:

    driver.c

Abstract:

    This file contains the driver entry points and callbacks.

Environment:

    Kernel-mode Driver Framework

--*/

#include "driver.h"
#include "driver.tmh"
#include <ntifs.h>
#include <wdm.h>
#include <ntdef.h>
#include "Structs.h"

NTKERNELAPI PVOID PsGetProcessSectionBaseAddress
(
	__in PEPROCESS Process
);

NTSTATUS NTAPI MmCopyVirtualMemory
(
	PEPROCESS SourceProcess,
	PVOID SourceAddress,
	PEPROCESS TargetProcess,
	PVOID TargetAddress,
	SIZE_T BufferSize,
	KPROCESSOR_MODE PreviousMode,
	PSIZE_T ReturnSize
);

PPEB
PsGetProcessPeb(
	__in PEPROCESS Process
);

#ifdef ALLOC_PRAGMA
#pragma alloc_text (INIT, DriverEntry)
#endif

ULONG64 llw;

PEPROCESS GetEPROCESS(int ProcessId) {
	PEPROCESS Process;
	return NT_SUCCESS(PsLookupProcessByProcessId((HANDLE)ProcessId, &Process)) ? Process : NULL;
}

ULONG64 GetProcAddress(ULONG64 block, LPCSTR name)
{
	PIMAGE_NT_HEADERS64 ntHeaders = (PIMAGE_NT_HEADERS64)((BYTE*)block + ((PIMAGE_DOS_HEADER)block)->e_lfanew);

	PIMAGE_DATA_DIRECTORY impDir = &ntHeaders->OptionalHeader.DataDirectory[0];
	PIMAGE_EXPORT_DIRECTORY ied = (PIMAGE_EXPORT_DIRECTORY)((BYTE*)block + impDir->VirtualAddress);

	for (DWORD64 i = 0; i < ied->NumberOfNames; i++)
	{
		DWORD64 curName = (DWORD64)(((BYTE*)block) + ied->AddressOfNames + i * sizeof(DWORD64));
		if (curName == 0 && strcmp(name, (LPSTR)((BYTE*)block + *(BYTE*)curName)) == 0)
		{
			unsigned short* pw = (unsigned short*)(((BYTE*)block) + ied->AddressOfNameOrdinals + i * sizeof(unsigned short));
			curName = (DWORD64)(((BYTE*)block) + ied->AddressOfFunctions + (*pw) * sizeof(DWORD64));
			return ((BYTE*)block + *(BYTE*)curName);
		}
	}
}

ULONG64 LoadLibrary(PEPROCESS target, PCWSTR baseName)
{
	UNICODE_STRING uName;
	OBJECT_ATTRIBUTES uAttributes;

	RtlInitUnicodeString(&uName, L"\\SystemRoot\\System32\\");
	RtlAppendUnicodeToString(&uName, baseName);

	InitializeObjectAttributes(&uAttributes, &uName, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);
	
	HANDLE   handle;
	IO_STATUS_BLOCK ioStatusBlock;
	IO_STATUS_BLOCK _ioStatusBlock;
	FILE_STANDARD_INFORMATION info;

	if (KeGetCurrentIrql() != PASSIVE_LEVEL)
		return 0;

	ZwCreateFile(&handle, GENERIC_WRITE, &uAttributes, &ioStatusBlock, NULL, FILE_ATTRIBUTE_NORMAL, 0, FILE_OVERWRITE_IF, FILE_SYNCHRONOUS_IO_NONALERT, NULL, 0);

	LARGE_INTEGER offset;
	offset.LowPart = offset.HighPart = 0;

	NtQueryInformationFile(handle, &_ioStatusBlock, &info, sizeof(FILE_STANDARD_INFORMATION), FileStandardInformation);

	PPEB peb = PsGetProcessPeb(target);
	ULONG totalCount = 0;

	PVOID buffer = ExAllocatePool(NonPagedPool, (SIZE_T)info.AllocationSize.QuadPart);
	ZwReadFile(handle, NULL, NULL, NULL, &ioStatusBlock, buffer, (SIZE_T)info.AllocationSize.QuadPart,&offset, NULL);

	PVOID allocationBase;
	SIZE_T rSize = wcslen(uName.Buffer);
	KAPC_STATE apcState;
	SIZE_T _size;
	HANDLE threadHandle;
	CLIENT_ID client;

	llw = GetProcAddress(buffer, "LoadLibraryW");

	for (PLIST_ENTRY pListEntry = peb->Ldr->InLoadOrderModuleList.Flink;
		pListEntry != &peb->Ldr->InLoadOrderModuleList;
		pListEntry = pListEntry->Flink)
	{
		totalCount++;
	}

	
	ULONG i = 0;
	for (PLIST_ENTRY pListEntry = peb->Ldr->InLoadOrderModuleList.Flink;
		pListEntry != &peb->Ldr->InLoadOrderModuleList;
		pListEntry = pListEntry->Flink)
	{
		PLDR_DATA_TABLE_ENTRY pEntry = CONTAINING_RECORD(pListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
		if (wcscmp(pEntry->BaseDllName, baseName) == 0)
		{
			return pEntry->DllBase;
		}
	}

	ZwAllocateVirtualMemory((HANDLE)-1, &allocationBase, 0, &rSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

	MmCopyVirtualMemory(IoGetCurrentProcess(), (PVOID)uName.Buffer, target, allocationBase, wcslen(uName.Buffer), KernelMode, &_size);

	PsCreateSystemThread(&threadHandle, THREAD_ALL_ACCESS, NULL, (HANDLE)-1, &client, (PKSTART_ROUTINE)GetProcAddress(buffer, "LoadLibraryW"), &allocationBase);
		
	for (PLIST_ENTRY pListEntry = peb->Ldr->InLoadOrderModuleList.Flink;
		pListEntry != &peb->Ldr->InLoadOrderModuleList;
		pListEntry = pListEntry->Flink)
	{
		totalCount++;
	}


	ULONG i = 0;
	ULONG j = 0;
	for (PLIST_ENTRY pListEntry = peb->Ldr->InLoadOrderModuleList.Flink;
		pListEntry != &peb->Ldr->InLoadOrderModuleList;
		pListEntry = pListEntry->Flink)
	{
		PLDR_DATA_TABLE_ENTRY pEntry = CONTAINING_RECORD(pListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
		if (wcscmp(pEntry->BaseDllName, baseName) == 0)
		{
			return pEntry->DllBase;
		}
	}

	return 0;
}

//Attach to process before use, inject() automatically does this
void ResolveImports(PEPROCESS target, PVOID image)
{
	DWORD64 i, Function, count, delta;

	ULONG64 hModule;
	PDWORD64 ptr;
	USHORT* list;

	PIMAGE_IMPORT_BY_NAME pIBN;
	PIMAGE_DOS_HEADER pIDH;	
	PIMAGE_THUNK_DATA64 FirstThunk, OrigFirstThunk;
	PIMAGE_NT_HEADERS64 pINH = (PIMAGE_NT_HEADERS)((BYTE*)image + pIDH->e_lfanew);
	PIMAGE_BASE_RELOCATION pIBR = (PIMAGE_BASE_RELOCATION)((BYTE*)image + pINH->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);;
 PIMAGE_IMPORT_DESCRIPTOR pIID = (PIMAGE_IMPORT_DESCRIPTOR)((BYTE*)image + pINH->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
	pIBR = (PIMAGE_BASE_RELOCATION)((BYTE*)image + pINH->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);

	while (pIBR->VirtualAddress)
	{
		if (pIBR->SizeOfBlock >= sizeof(IMAGE_BASE_RELOCATION))
		{
			count = (pIBR->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(USHORT);
			list = (USHORT*)(pIBR + 1);

			for (i = 0; i < count; i++)
			{
				if (list[i])
				{
					ptr = (DWORD64)((BYTE*)image + (pIBR->VirtualAddress + (list[i] & 0xFFF)));
					*ptr += delta;
				}
			}
		}

		pIBR = (PIMAGE_BASE_RELOCATION)((BYTE*)pIBR + pIBR->SizeOfBlock);
	}

	pIID = pIID;

	
	while (pIID->Characteristics)
	{
		OrigFirstThunk = (PIMAGE_THUNK_DATA64)((BYTE*)image + pIID->OriginalFirstThunk);
		FirstThunk = (PIMAGE_THUNK_DATA64)((BYTE*)image + pIID->FirstThunk);

		hModule = LoadLibrary(target, pIID->Name);

		if (!hModule)
		{
			return FALSE;
		}

		while (OrigFirstThunk->u1.AddressOfData)
		{
			if (OrigFirstThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG64)
			{
				// Import by ordinal

				Function = (DWORD64)GetProcAddress(hModule, (LPCSTR)(OrigFirstThunk->u1.Ordinal & 0xFFFF));

				if (!Function)
				{
					return FALSE;
				}

				FirstThunk->u1.Function = Function;
			}

			else
			{
				// Import by name

				pIBN = (PIMAGE_IMPORT_BY_NAME)((BYTE*)image + OrigFirstThunk->u1.AddressOfData);
				Function = (DWORD64)GetProcAddress(hModule, (LPCSTR)pIBN->Name);

				if (!Function)
				{
					return FALSE;
				}

				FirstThunk->u1.Function = Function;
			}

			OrigFirstThunk++;
			FirstThunk++;
		}

		pIID++;
	}
}

void inject(PEPROCESS target, void* image, ULONG64 size, ULONG64 entry)
{
	KAPC_STATE apcState;
	ULONG64 allocationBase;
	ULONG64 rSize;
	HANDLE thread;
	CLIENT_ID client;

	KeStackAttachProcess(target, &apcState);
	ResolveImports(target, image);
	ZwAllocateVirtualMemory((HANDLE)-1, &allocationBase, 0, &rSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
	memcpy(allocationBase, image, rSize);
	PsCreateSystemThread(&thread, THREAD_ALL_ACCESS, NULL, (HANDLE)-1, &client, (PKSTART_ROUTINE)entry, NULL);
	KeUnstackDetachProcess(&apcState);
}
NTSTATUS
DriverEntry(
    _In_ PDRIVER_OBJECT  DriverObject,
    _In_ PUNICODE_STRING RegistryPath
    )
{
	
	NTSTATUS NtStatus = STATUS_SUCCESS;
	return NtStatus;
}