noury/mirror-ac
1
0
Fork 0
mirror of https://github.com/donnaskiez/ac.git synced 2024-11-21 22:24:08 +01:00
Code Issues Projects Releases Packages Wiki Activity

refactor some more

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

88
driver/integrity.c
View file

@ -8,6 +8,7 @@
#include "imports.h"
#include "session.h"
#include "util.h"
#include "pe.h"
#include <bcrypt.h>
#include <initguid.h>
@ -205,6 +206,33 @@ IsSectionExecutable(_In_ PIMAGE_SECTION_HEADER Section)
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
NTSTATUS
StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
@ -217,7 +245,7 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
NTSTATUS status = STATUS_UNSUCCESSFUL;
PIMAGE_DOS_HEADER dos_header = NULL;
PLOCAL_NT_HEADER nt_header = NULL;
PNT_HEADER_64 nt_header = NULL;
PIMAGE_SECTION_HEADER section = NULL;
ULONG total_packet_size = 0;
ULONG num_sections = 0;
@ -225,11 +253,14 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
UINT64 buffer_base = 0;
ULONG bytes_returned = 0;
MM_COPY_ADDRESS address = {0};
ULONG buffer_size = 0;
INTEGRITY_CHECK_HEADER header = {0};
if (!ModuleBase || !ModuleSize)
return STATUS_INVALID_PARAMETER;
if (!IsModuleAddressSafe(ModuleBase, IsModulex86))
return STATUS_UNSUCCESSFUL;
/*
* 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
@ -237,37 +268,20 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
* our reponse until we enumerate and count all executable sections for
* the file.
*/
buffer_size = ModuleSize + sizeof(INTEGRITY_CHECK_HEADER);
*BytesWritten = 0;
*Buffer = ImpExAllocatePool2(
POOL_FLAG_NON_PAGED, buffer_size, POOL_TAG_INTEGRITY);
*Buffer = ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
ModuleSize + sizeof(INTEGRITY_CHECK_HEADER),
POOL_TAG_INTEGRITY);
if (*Buffer == NULL)
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
* IMAGE_NT_HEADER from the base of the image.
*/
nt_header = (struct _IMAGE_NT_HEADERS64*)((UINT64)ModuleBase +
dos_header->e_lfanew);
num_sections = nt_header->FileHeader.NumberOfSections;
nt_header = PeGetNtHeader(ModuleBase);
num_sections = GetSectionCount(nt_header);
/*
* The IMAGE_FIRST_SECTION macro takes in an IMAGE_NT_HEADER and returns
@ -283,7 +297,6 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
}
address.VirtualAddress = section;
status = ImpMmCopyMemory((UINT64)buffer_base + total_packet_size,
address,
sizeof(IMAGE_SECTION_HEADER),
@ -291,14 +304,12 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
&bytes_returned);
if (!NT_SUCCESS(status)) {
DEBUG_ERROR("MmCopyMemory failed with status %x", status);
ImpExFreePoolWithTag(*Buffer, POOL_TAG_INTEGRITY);
*Buffer = NULL;
return status;
}
address.VirtualAddress = (UINT64)ModuleBase + section->PointerToRawData;
status = ImpMmCopyMemory((UINT64)buffer_base + total_packet_size +
sizeof(IMAGE_SECTION_HEADER),
address,
@ -307,22 +318,18 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
&bytes_returned);
if (!NT_SUCCESS(status)) {
DEBUG_ERROR("MmCopyMemory failed with status %x", status);
ImpExFreePoolWithTag(*Buffer, POOL_TAG_INTEGRITY);
*Buffer = NULL;
return status;
}
total_packet_size +=
section->SizeOfRawData + sizeof(IMAGE_SECTION_HEADER);
total_packet_size += GetSectionTotalPacketSize(section);
num_executable_sections++;
section++;
}
INTEGRITY_CHECK_HEADER header = {0};
header.executable_section_count = num_executable_sections;
header.total_packet_size =
total_packet_size + sizeof(INTEGRITY_CHECK_HEADER);
InitIntegrityCheckHeader(
&header, num_executable_sections, total_packet_size);
RtlCopyMemory(*Buffer, &header, sizeof(INTEGRITY_CHECK_HEADER));
*BytesWritten = total_packet_size + sizeof(INTEGRITY_CHECK_HEADER);
@ -332,10 +339,9 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer,
STATIC
NTSTATUS
MapDiskImageIntoVirtualAddressSpace(_Inout_ PHANDLE SectionHandle,
_Outptr_result_bytebuffer_(*Size)
PVOID* Section,
_Out_ PVOID* Section,
_In_ PUNICODE_STRING Path,
_Out_ _Deref_out_range_(>, 0) PSIZE_T Size)
_Out_ PSIZE_T Size)
{
PAGED_CODE();
@ -1197,20 +1203,18 @@ UINT64
MeasureReads(_In_ PVOID Address, _In_ ULONG Count)
{
UINT64 read_average = 0;
UINT64 old_irql = 0;
KIRQL irql = {0};
MeasureInstructionRead(Address);
old_irql = __readcr8();
__writecr8(HIGH_LEVEL);
KeRaiseIrql(HIGH_LEVEL, &irql);
_disable();
for (ULONG iteration = 0; iteration < Count; iteration++)
read_average += MeasureInstructionRead(Address);
_enable();
__writecr8(old_irql);
KeLowerIrql(irql);
DEBUG_VERBOSE("EPT Detection - Read Average: %llx", read_average);

61
driver/modules.c
View file

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

6
driver/pe.c
View file

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

13
driver/pe.h
View file

@ -9,4 +9,17 @@
PVOID
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