Domito/src/Domito.cpp

#include <ntifs.h>
#include <ntimage.h>

#include "Domito.h"


// Structure representing a loaded module
typedef struct _SYSTEM_MODULE_INFORMATION_ENTRY
{
	PVOID Unknown1;
	PVOID Unknown2;
	PVOID Base;
	ULONG Size;
	ULONG Flags;
	USHORT Index;
	USHORT NameLength;
	USHORT LoadCount;
	USHORT PathLength;
	CHAR ImageName[256];
} SYSTEM_MODULE_INFORMATION_ENTRY, * PSYSTEM_MODULE_INFORMATION_ENTRY;

// Structure representing the loaded module information
typedef struct _SYSTEM_MODULE_INFORMATION
{
	ULONG Count;
	SYSTEM_MODULE_INFORMATION_ENTRY Module[1];
} SYSTEM_MODULE_INFORMATION, * PSYSTEM_MODULE_INFORMATION;

// Function prototype for ZwQuerySystemInformation
NTSYSAPI NTSTATUS NTAPI ZwQuerySystemInformation(
	ULONG SystemInformationClass,
	PVOID SystemInformation,
	ULONG SystemInformationLength,
	PULONG ReturnLength
);

typedef struct _LDR_DATA_TABLE_ENTRY
{
	LIST_ENTRY64 InLoadOrderLinks;
	PVOID ExceptionTable;
	ULONG ExceptionTableSize;
	PVOID GpValue;
	PVOID NonPagedDebugInfo;
	PVOID ImageBase;
	PVOID EntryPoint;
	ULONG SizeOfImage;
	UNICODE_STRING FullImageName;
	UNICODE_STRING BaseImageName;
	ULONG Flags;
	USHORT LoadCount;
	USHORT TlsIndex;
	LIST_ENTRY64 HashLinks;
	PVOID SectionPointer;
	ULONG CheckSum;
	ULONG TimeDateStamp;
	PVOID LoadedImports;
	PVOID EntryPointActivationContext;
	PVOID PatchInformation;
} LDR_DATA_TABLE_ENTRY, * PLDR_DATA_TABLE_ENTRY;

typedef PVOID(NTAPI* t_RtlImageDirectoryEntryToData)(
	IN PVOID Base,
	IN BOOLEAN MappedAsImage,
	IN USHORT DirectoryEntry,
	OUT PULONG Size
	);


_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoFindDriverBaseAddress(
	_In_ STRING ModuleName,
	_In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
	_Inout_opt_ PVOID * ModuleBase
)
{
	ULONG bufferSize = 0;
	PSYSTEM_MODULE_INFORMATION moduleInfo = NULL;

	const ULONG SystemModuleInformation = 11;

	// Query the required buffer size for module information
	NTSTATUS status = ZwQuerySystemInformation(
		SystemModuleInformation,
		&bufferSize,
		0,
		&bufferSize
	);

	if (status != STATUS_INFO_LENGTH_MISMATCH)
	{
		return status;
	}

#pragma warning(disable:4996)
	// Allocate memory for the module information
	moduleInfo = (PSYSTEM_MODULE_INFORMATION)Allocator(
		bufferSize
	);
#pragma warning(default:4996)

	if (moduleInfo == NULL)
	{
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	// Retrieve the module information
	status = ZwQuerySystemInformation(
		SystemModuleInformation,
		moduleInfo,
		bufferSize,
		NULL
	);

	if (!NT_SUCCESS(status))
	{
		ExFreePool(moduleInfo);
		return status;
	}

	STRING currentImageName;

	status = STATUS_NOT_FOUND;
	// Iterate through the loaded modules and find the desired module
	for (ULONG i = 0; i < moduleInfo->Count; i++)
	{
		RtlInitAnsiString(&currentImageName, moduleInfo->Module[i].ImageName);

		if (0 == RtlCompareString(&ModuleName, &currentImageName, TRUE))
		{
			// Found the module, store the base address
			if (ModuleBase)
			{
				status = STATUS_SUCCESS;
				*ModuleBase = moduleInfo->Module[i].Base;
			}
			break;
		}
	}

	ExFreePool(moduleInfo);

	return status;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoFindExportedFunctionAddress(
	_In_ PVOID ModuleBase,
	_In_ STRING FunctionName,
	_Inout_opt_ PVOID * FunctionAddress
)
{
	NTSTATUS status = STATUS_NOT_FOUND;
	ULONG exportSize;

	DECLARE_CONST_UNICODE_STRING(routineName, L"RtlImageDirectoryEntryToData");

	const t_RtlImageDirectoryEntryToData fp_RtlImageDirectoryEntryToData =
		(t_RtlImageDirectoryEntryToData)MmGetSystemRoutineAddress((PUNICODE_STRING)&routineName);

	if (fp_RtlImageDirectoryEntryToData == NULL)
	{
		return STATUS_NOT_IMPLEMENTED;
	}

	// Retrieve the export directory information
	const PIMAGE_EXPORT_DIRECTORY exportDirectory = (PIMAGE_EXPORT_DIRECTORY)fp_RtlImageDirectoryEntryToData(
		ModuleBase,
		TRUE,
		IMAGE_DIRECTORY_ENTRY_EXPORT,
		&exportSize
	);

	if (exportDirectory == NULL)
	{
		return STATUS_INVALID_IMAGE_FORMAT;
	}

	STRING currentFunctionName;

	const PULONG functionAddresses = (PULONG)((ULONG_PTR)ModuleBase + exportDirectory->AddressOfFunctions);
	const PULONG functionNames = (PULONG)((ULONG_PTR)ModuleBase + exportDirectory->AddressOfNames);
	const PUSHORT functionOrdinals = (PUSHORT)((ULONG_PTR)ModuleBase + exportDirectory->AddressOfNameOrdinals);

	for (ULONG i = 0; i < exportDirectory->NumberOfNames; i++)
	{
		const char* functionName = (const char*)((ULONG_PTR)ModuleBase + functionNames[i]);
		const USHORT functionOrdinal = functionOrdinals[i];
		UNREFERENCED_PARAMETER(functionOrdinal);

		const ULONG functionRva = functionAddresses[i];
		const PVOID functionAddress = (PVOID)((ULONG_PTR)ModuleBase + functionRva);

		RtlInitAnsiString(&currentFunctionName, functionName);

		if (0 == RtlCompareString(&FunctionName, &currentFunctionName, TRUE))
		{
			if (FunctionAddress)
			{
				status = STATUS_SUCCESS;
				*FunctionAddress = functionAddress;
			}
			break;
		}
	}

	return status;
}