#include "pch.h"

#include "KernelModuleTracker.h"

#include <PEParser.h>

#include "Helpers.h"

using namespace WinSys;

uint32_t WinSys::KernelModuleTracker::EnumModules() {

_modules.clear();

DWORD size = 1 << 18;

wil::unique_virtualalloc_ptr<> buffer(::VirtualAlloc(nullptr, size,MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE));

PVOID SystemRangeStart = 0;

NtQuerySystemInformation(SystemRangeStartInformation,

&SystemRangeStart,

sizeof(SystemRangeStart),

NULL);

NTSTATUS status;

status = ::NtQuerySystemInformation(SystemModuleInformationEx, buffer.get(), size, nullptr);

if (!NT_SUCCESS(status))

return 0;

if (_modules.empty()) {

_modules.reserve(256);

}

auto p = (RTL_PROCESS_MODULE_INFORMATION_EX*)buffer.get();

CHAR winDir[MAX_PATH];

::GetWindowsDirectoryA(winDir, _countof(winDir));

static const std::string root("\\SystemRoot\\");

static const std::string global("\\??\\");

for (;;) {

if (p->BaseInfo.ImageBase == 0)

break;

auto m = std::make_shared<KernelModuleInfo>();

m->Flags = p->BaseInfo.Flags;

m->FullPath = (const char*)p->BaseInfo.FullPathName;

m->NtPath = m->FullPath;

if (m->FullPath.find(global) == 0)

m->FullPath = m->FullPath.substr(global.size());

if (m->FullPath.find(root) == 0)

m->FullPath = winDir + m->FullPath.substr(root.size() - 1);

if (p->BaseInfo.ImageBase < SystemRangeStart) {

if (p->NextOffset == 0)

break;

p = (RTL_PROCESS_MODULE_INFORMATION_EX*)((BYTE*)p + p->NextOffset);

continue;

}

std::wstring path = Helpers::StringToWstring(m->FullPath);

PEParser parser(path.c_str());

if (parser.IsValid()) {

auto type = parser.GetSubsystemType();

if (type != SubsystemType::Native) {

if (p->NextOffset == 0)

break;

p = (RTL_PROCESS_MODULE_INFORMATION_EX*)((BYTE*)p + p->NextOffset);

continue;

}

}

m->ImageBase = p->BaseInfo.ImageBase;

m->MappedBase = p->BaseInfo.MappedBase;

m->ImageSize = p->BaseInfo.ImageSize;

m->InitOrderIndex = p->BaseInfo.InitOrderIndex;

m->LoadOrderIndex = p->BaseInfo.LoadOrderIndex;

m->LoadCount = p->BaseInfo.LoadCount;

m->hSection = p->BaseInfo.Section;

m->DefaultBase = p->DefaultBase;

m->ImageChecksum = p->ImageChecksum;

m->TimeDateStamp = p->TimeDateStamp;

m->Name = std::string((PCSTR)((BYTE*)p->BaseInfo.FullPathName + p->BaseInfo.OffsetToFileName));

_modules.push_back(std::move(m));

if (p->NextOffset == 0)

break;

p = (RTL_PROCESS_MODULE_INFORMATION_EX*)((BYTE*)p + p->NextOffset);

}

return uint32_t(_modules.size());

}

const std::vector<std::shared_ptr<KernelModuleInfo>>& WinSys::KernelModuleTracker::GetModules() const {

return _modules;

}