refactor some more

This commit is contained in:
donnaskiez 2024-05-05 23:20:35 +10:00
parent 80b34ea2f1
commit 8685725415
4 changed files with 103 additions and 103 deletions

View file

@ -8,6 +8,7 @@
#include "imports.h" #include "imports.h"
#include "session.h" #include "session.h"
#include "util.h" #include "util.h"
#include "pe.h"
#include <bcrypt.h> #include <bcrypt.h>
#include <initguid.h> #include <initguid.h>
@ -205,6 +206,33 @@ IsSectionExecutable(_In_ PIMAGE_SECTION_HEADER Section)
return Section->Characteristics & IMAGE_SCN_MEM_EXECUTE ? TRUE : FALSE; return Section->Characteristics & IMAGE_SCN_MEM_EXECUTE ? TRUE : FALSE;
} }
FORCEINLINE
STATIC
BOOLEAN
IsModuleAddressSafe(_In_ PVOID Base, _In_ BOOLEAN x86)
{
return !MmIsAddressValid(Base) && !x86 ? FALSE : TRUE;
}
FORCEINLINE
STATIC
UINT32
GetSectionTotalPacketSize(_In_ PIMAGE_SECTION_HEADER Section)
{
return Section->SizeOfRawData + sizeof(IMAGE_SECTION_HEADER);
}
FORCEINLINE
STATIC
VOID
InitIntegrityCheckHeader(_Out_ PINTEGRITY_CHECK_HEADER Header,
_In_ UINT32 SectionCount,
_In_ UINT32 TotalSize)
{
Header->executable_section_count = SectionCount;
Header->total_packet_size = TotalSize + sizeof(INTEGRITY_CHECK_HEADER);
}
STATIC STATIC
NTSTATUS NTSTATUS
StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer, StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
@ -215,21 +243,24 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
{ {
PAGED_CODE(); PAGED_CODE();
NTSTATUS status = STATUS_UNSUCCESSFUL; NTSTATUS status = STATUS_UNSUCCESSFUL;
PIMAGE_DOS_HEADER dos_header = NULL; PIMAGE_DOS_HEADER dos_header = NULL;
PLOCAL_NT_HEADER nt_header = NULL; PNT_HEADER_64 nt_header = NULL;
PIMAGE_SECTION_HEADER section = NULL; PIMAGE_SECTION_HEADER section = NULL;
ULONG total_packet_size = 0; ULONG total_packet_size = 0;
ULONG num_sections = 0; ULONG num_sections = 0;
ULONG num_executable_sections = 0; ULONG num_executable_sections = 0;
UINT64 buffer_base = 0; UINT64 buffer_base = 0;
ULONG bytes_returned = 0; ULONG bytes_returned = 0;
MM_COPY_ADDRESS address = {0}; MM_COPY_ADDRESS address = {0};
ULONG buffer_size = 0; INTEGRITY_CHECK_HEADER header = {0};
if (!ModuleBase || !ModuleSize) if (!ModuleBase || !ModuleSize)
return STATUS_INVALID_PARAMETER; return STATUS_INVALID_PARAMETER;
if (!IsModuleAddressSafe(ModuleBase, IsModulex86))
return STATUS_UNSUCCESSFUL;
/* /*
* The reason we allocate a buffer to temporarily hold the section data * The reason we allocate a buffer to temporarily hold the section data
* is that we don't know the total size until after we iterate over the * is that we don't know the total size until after we iterate over the
@ -237,37 +268,20 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
* our reponse until we enumerate and count all executable sections for * our reponse until we enumerate and count all executable sections for
* the file. * the file.
*/ */
buffer_size = ModuleSize + sizeof(INTEGRITY_CHECK_HEADER);
*BytesWritten = 0; *BytesWritten = 0;
*Buffer = ImpExAllocatePool2( *Buffer = ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
POOL_FLAG_NON_PAGED, buffer_size, POOL_TAG_INTEGRITY); ModuleSize + sizeof(INTEGRITY_CHECK_HEADER),
POOL_TAG_INTEGRITY);
if (*Buffer == NULL) if (*Buffer == NULL)
return STATUS_MEMORY_NOT_ALLOCATED; return STATUS_MEMORY_NOT_ALLOCATED;
/*
* Note: Verifier doesn't like it when we map the module so rather then
* mapping it to our address space we will simply use MmCopyMemory on
* the module to avoid upsetting verifier
* :)
*/
dos_header = (PIMAGE_DOS_HEADER)ModuleBase;
if (!MmIsAddressValid(dos_header) && !IsModulex86) {
ImpExFreePoolWithTag(*Buffer, POOL_TAG_INTEGRITY);
*Buffer = NULL;
return STATUS_INVALID_ADDRESS;
}
/* /*
* The IMAGE_DOS_HEADER.e_lfanew stores the offset of the * The IMAGE_DOS_HEADER.e_lfanew stores the offset of the
* IMAGE_NT_HEADER from the base of the image. * IMAGE_NT_HEADER from the base of the image.
*/ */
nt_header = (struct _IMAGE_NT_HEADERS64*)((UINT64)ModuleBase + nt_header = PeGetNtHeader(ModuleBase);
dos_header->e_lfanew); num_sections = GetSectionCount(nt_header);
num_sections = nt_header->FileHeader.NumberOfSections;
/* /*
* The IMAGE_FIRST_SECTION macro takes in an IMAGE_NT_HEADER and returns * The IMAGE_FIRST_SECTION macro takes in an IMAGE_NT_HEADER and returns
@ -283,7 +297,6 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
} }
address.VirtualAddress = section; address.VirtualAddress = section;
status = ImpMmCopyMemory((UINT64)buffer_base + total_packet_size, status = ImpMmCopyMemory((UINT64)buffer_base + total_packet_size,
address, address,
sizeof(IMAGE_SECTION_HEADER), sizeof(IMAGE_SECTION_HEADER),
@ -291,14 +304,12 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
&bytes_returned); &bytes_returned);
if (!NT_SUCCESS(status)) { if (!NT_SUCCESS(status)) {
DEBUG_ERROR("MmCopyMemory failed with status %x", status);
ImpExFreePoolWithTag(*Buffer, POOL_TAG_INTEGRITY); ImpExFreePoolWithTag(*Buffer, POOL_TAG_INTEGRITY);
*Buffer = NULL; *Buffer = NULL;
return status; return status;
} }
address.VirtualAddress = (UINT64)ModuleBase + section->PointerToRawData; address.VirtualAddress = (UINT64)ModuleBase + section->PointerToRawData;
status = ImpMmCopyMemory((UINT64)buffer_base + total_packet_size + status = ImpMmCopyMemory((UINT64)buffer_base + total_packet_size +
sizeof(IMAGE_SECTION_HEADER), sizeof(IMAGE_SECTION_HEADER),
address, address,
@ -307,22 +318,18 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
&bytes_returned); &bytes_returned);
if (!NT_SUCCESS(status)) { if (!NT_SUCCESS(status)) {
DEBUG_ERROR("MmCopyMemory failed with status %x", status);
ImpExFreePoolWithTag(*Buffer, POOL_TAG_INTEGRITY); ImpExFreePoolWithTag(*Buffer, POOL_TAG_INTEGRITY);
*Buffer = NULL; *Buffer = NULL;
return status; return status;
} }
total_packet_size += total_packet_size += GetSectionTotalPacketSize(section);
section->SizeOfRawData + sizeof(IMAGE_SECTION_HEADER);
num_executable_sections++; num_executable_sections++;
section++; section++;
} }
INTEGRITY_CHECK_HEADER header = {0}; InitIntegrityCheckHeader(
header.executable_section_count = num_executable_sections; &header, num_executable_sections, total_packet_size);
header.total_packet_size =
total_packet_size + sizeof(INTEGRITY_CHECK_HEADER);
RtlCopyMemory(*Buffer, &header, sizeof(INTEGRITY_CHECK_HEADER)); RtlCopyMemory(*Buffer, &header, sizeof(INTEGRITY_CHECK_HEADER));
*BytesWritten = total_packet_size + sizeof(INTEGRITY_CHECK_HEADER); *BytesWritten = total_packet_size + sizeof(INTEGRITY_CHECK_HEADER);
@ -331,11 +338,10 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
STATIC STATIC
NTSTATUS NTSTATUS
MapDiskImageIntoVirtualAddressSpace(_Inout_ PHANDLE SectionHandle, MapDiskImageIntoVirtualAddressSpace(_Inout_ PHANDLE SectionHandle,
_Outptr_result_bytebuffer_(*Size) _Out_ PVOID* Section,
PVOID* Section,
_In_ PUNICODE_STRING Path, _In_ PUNICODE_STRING Path,
_Out_ _Deref_out_range_(>, 0) PSIZE_T Size) _Out_ PSIZE_T Size)
{ {
PAGED_CODE(); PAGED_CODE();
@ -1197,20 +1203,18 @@ UINT64
MeasureReads(_In_ PVOID Address, _In_ ULONG Count) MeasureReads(_In_ PVOID Address, _In_ ULONG Count)
{ {
UINT64 read_average = 0; UINT64 read_average = 0;
UINT64 old_irql = 0; KIRQL irql = {0};
MeasureInstructionRead(Address); MeasureInstructionRead(Address);
old_irql = __readcr8(); KeRaiseIrql(HIGH_LEVEL, &irql);
__writecr8(HIGH_LEVEL);
_disable(); _disable();
for (ULONG iteration = 0; iteration < Count; iteration++) for (ULONG iteration = 0; iteration < Count; iteration++)
read_average += MeasureInstructionRead(Address); read_average += MeasureInstructionRead(Address);
_enable(); _enable();
__writecr8(old_irql); KeLowerIrql(irql);
DEBUG_VERBOSE("EPT Detection - Read Average: %llx", read_average); DEBUG_VERBOSE("EPT Detection - Read Average: %llx", read_average);

View file

@ -75,9 +75,9 @@ typedef struct _NMI_CONTEXT {
} NMI_CONTEXT, *PNMI_CONTEXT; } NMI_CONTEXT, *PNMI_CONTEXT;
STATIC STATIC
NTSTATUS VOID
PopulateWhitelistedModuleBuffer(_Inout_ PVOID Buffer, PopulateWhitelistedModuleBuffer(_Inout_ PWHITELISTED_REGIONS Whitelist,
_In_ PSYSTEM_MODULES SystemModules); _In_ PSYSTEM_MODULES SystemModules);
STATIC STATIC
NTSTATUS NTSTATUS
@ -160,48 +160,27 @@ FindSystemModuleByName(_In_ LPCSTR ModuleName,
return NULL; return NULL;
} }
FORCEINLINE
STATIC STATIC
VOID VOID
InitWhitelistedRegionStructure(_Out_ PWHITELISTED_REGIONS Region, PopulateWhitelistedModuleBuffer(_Inout_ PWHITELISTED_REGIONS Whitelist,
_In_ UINT64 Base, _In_ PSYSTEM_MODULES SystemModules)
_In_ UINT64 End)
{
Region->base = Base;
Region->end = End;
}
STATIC
NTSTATUS
PopulateWhitelistedModuleBuffer(_Inout_ PVOID Buffer,
_In_ PSYSTEM_MODULES SystemModules)
{ {
PAGED_CODE(); PAGED_CODE();
if (!Buffer || !SystemModules)
return STATUS_INVALID_PARAMETER;
for (INT index = 0; index < WHITELISTED_MODULE_COUNT; index++) { for (INT index = 0; index < WHITELISTED_MODULE_COUNT; index++) {
LPCSTR name = WHITELISTED_MODULES[index]; LPCSTR entry = WHITELISTED_MODULES[index];
PRTL_MODULE_EXTENDED_INFO module = PRTL_MODULE_EXTENDED_INFO module =
FindSystemModuleByName(name, SystemModules); FindSystemModuleByName(entry, SystemModules);
/* not everyone will contain all whitelisted modules */ /* not everyone will contain all whitelisted modules */
if (!module) if (!module)
continue; continue;
WHITELISTED_REGIONS region = {0}; PWHITELISTED_REGIONS region = &Whitelist[index];
InitWhitelistedRegionStructure(&region, region->base = (UINT64)module->ImageBase;
(UINT64)module->ImageBase, region->end = (UINT64)module->ImageBase + module->ImageSize;
region.base + module->ImageSize);
UINT64 destination =
(UINT64)Buffer + index * sizeof(WHITELISTED_REGIONS);
RtlCopyMemory(destination, &region, sizeof(WHITELISTED_REGIONS));
} }
return STATUS_SUCCESS;
} }
STATIC STATIC
@ -400,7 +379,7 @@ STATIC
VOID VOID
ValidateDriverObjects(_In_ PSYSTEM_MODULES Modules, ValidateDriverObjects(_In_ PSYSTEM_MODULES Modules,
_In_ POBJECT_DIRECTORY_ENTRY Entry, _In_ POBJECT_DIRECTORY_ENTRY Entry,
_In_ PVOID Whitelist) _In_ PWHITELISTED_REGIONS Whitelist)
{ {
NTSTATUS status = STATUS_UNSUCCESSFUL; NTSTATUS status = STATUS_UNSUCCESSFUL;
POBJECT_DIRECTORY_ENTRY entry = Entry; POBJECT_DIRECTORY_ENTRY entry = Entry;
@ -412,7 +391,7 @@ ValidateDriverObjects(_In_ PSYSTEM_MODULES Modules,
ReportInvalidDriverObject(driver, REPORT_SUBTYPE_NO_BACKING_MODULE); ReportInvalidDriverObject(driver, REPORT_SUBTYPE_NO_BACKING_MODULE);
} }
if (!DoesDriverHaveInvalidDispatchRoutine(driver, Modules, Whitelist)) { if (DoesDriverHaveInvalidDispatchRoutine(driver, Modules, Whitelist)) {
ReportInvalidDriverObject(driver, REPORT_SUBTYPE_INVALID_DISPATCH); ReportInvalidDriverObject(driver, REPORT_SUBTYPE_INVALID_DISPATCH);
} }
@ -430,13 +409,13 @@ ValidateDriverObjectsWrapper(_In_ PSYSTEM_MODULES SystemModules)
{ {
PAGED_CODE(); PAGED_CODE();
HANDLE handle = NULL; HANDLE handle = NULL;
OBJECT_ATTRIBUTES attributes = {0}; OBJECT_ATTRIBUTES attributes = {0};
PVOID directory = {0}; PVOID directory = {0};
UNICODE_STRING directory_name = {0}; UNICODE_STRING directory_name = {0};
PVOID whitelisted_regions_buffer = NULL; PWHITELISTED_REGIONS whitelist = NULL;
NTSTATUS status = STATUS_UNSUCCESSFUL; NTSTATUS status = STATUS_UNSUCCESSFUL;
POBJECT_DIRECTORY directory_object = NULL; POBJECT_DIRECTORY directory_object = NULL;
ImpRtlInitUnicodeString(&directory_name, L"\\Driver"); ImpRtlInitUnicodeString(&directory_name, L"\\Driver");
@ -477,16 +456,15 @@ ValidateDriverObjectsWrapper(_In_ PSYSTEM_MODULES SystemModules)
ImpExAcquirePushLockExclusiveEx(&directory_object->Lock, NULL); ImpExAcquirePushLockExclusiveEx(&directory_object->Lock, NULL);
whitelisted_regions_buffer = whitelist = ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
ImpExAllocatePool2(POOL_FLAG_NON_PAGED, WHITELISTED_MODULE_COUNT *
WHITELISTED_MODULE_COUNT * MODULE_MAX_STRING_SIZE, sizeof(WHITELISTED_REGIONS),
WHITELISTED_MODULE_TAG); WHITELISTED_MODULE_TAG);
if (!whitelisted_regions_buffer) if (!whitelist)
goto end; goto end;
status = PopulateWhitelistedModuleBuffer(whitelisted_regions_buffer, PopulateWhitelistedModuleBuffer(whitelist, SystemModules);
SystemModules);
if (!NT_SUCCESS(status)) { if (!NT_SUCCESS(status)) {
DEBUG_ERROR("PopulateWhitelistedModuleBuffer failed with status %x", DEBUG_ERROR("PopulateWhitelistedModuleBuffer failed with status %x",
@ -496,13 +474,12 @@ ValidateDriverObjectsWrapper(_In_ PSYSTEM_MODULES SystemModules)
for (INT index = 0; index < NUMBER_HASH_BUCKETS; index++) { for (INT index = 0; index < NUMBER_HASH_BUCKETS; index++) {
POBJECT_DIRECTORY_ENTRY entry = directory_object->HashBuckets[index]; POBJECT_DIRECTORY_ENTRY entry = directory_object->HashBuckets[index];
ValidateDriverObjects(SystemModules, entry, whitelisted_regions_buffer); ValidateDriverObjects(SystemModules, entry, whitelist);
} }
end: end:
if (whitelisted_regions_buffer) if (whitelist)
ImpExFreePoolWithTag(whitelisted_regions_buffer, ImpExFreePoolWithTag(whitelist, WHITELISTED_MODULE_TAG);
WHITELISTED_MODULE_TAG);
ImpExReleasePushLockExclusiveEx(&directory_object->Lock, 0); ImpExReleasePushLockExclusiveEx(&directory_object->Lock, 0);
ImpObDereferenceObject(directory); ImpObDereferenceObject(directory);

View file

@ -34,6 +34,12 @@ PeGetExportDirectory(_In_ PVOID Image,
PIMAGE_EXPORT_DIRECTORY, Image, ExportDataDirectory->VirtualAddress); PIMAGE_EXPORT_DIRECTORY, Image, ExportDataDirectory->VirtualAddress);
} }
UINT32
GetSectionCount(_In_ PNT_HEADER_64 Header)
{
return Header->FileHeader.NumberOfSections;
}
PVOID PVOID
PeFindExportByName(_In_ PVOID Image, _In_ PCHAR Name) PeFindExportByName(_In_ PVOID Image, _In_ PCHAR Name)
{ {

View file

@ -9,4 +9,17 @@
PVOID PVOID
PeFindExportByName(_In_ PVOID Image, _In_ PCHAR Name); PeFindExportByName(_In_ PVOID Image, _In_ PCHAR Name);
PNT_HEADER_64
PeGetNtHeader(_In_ PVOID Image);
PIMAGE_DATA_DIRECTORY
PeGetExportDataDirectory(_In_ PVOID Image);
PIMAGE_EXPORT_DIRECTORY
PeGetExportDirectory(_In_ PVOID Image,
_In_ PIMAGE_DATA_DIRECTORY ExportDataDirectory);
UINT32
GetSectionCount(_In_ PNT_HEADER_64 Header);
#endif #endif