mirror of
https://github.com/hasherezade/pe-sieve
synced 2026-06-08 14:34:52 +00:00
hasherezade
109 lines
3.5 KiB
C++
109 lines
3.5 KiB
C++
#include "workingset_enum.h"
|
|
|
|
#include <iostream>
|
|
|
|
#include <psapi.h>
|
|
#pragma comment(lib,"psapi.lib")
|
|
|
|
#ifdef _WIN64
|
|
const ULONGLONG mask = ULONGLONG(-1);
|
|
#else
|
|
const ULONGLONG mask = DWORD(-1);
|
|
#endif
|
|
|
|
namespace pesieve {
|
|
namespace util {
|
|
|
|
bool get_next_commited_region(HANDLE processHandle, ULONGLONG start_va, MEMORY_BASIC_INFORMATION &page_info)
|
|
{
|
|
const SIZE_T page_info_size = sizeof(MEMORY_BASIC_INFORMATION);
|
|
|
|
while (start_va < mask) {
|
|
//std::cout << "Checking: " << std::hex << start_va << std::endl;
|
|
memset(&page_info, 0, page_info_size);
|
|
const SIZE_T out = VirtualQueryEx(processHandle, (LPCVOID)start_va, &page_info, page_info_size);
|
|
const bool is_read = (out == page_info_size) ? true : false;
|
|
const DWORD error = is_read ? ERROR_SUCCESS : GetLastError();
|
|
if (error == ERROR_INVALID_PARAMETER) {
|
|
//nothing more to read
|
|
#ifdef _DEBUG
|
|
std::cout << "Nothing more to read: " << std::hex << start_va << std::endl;
|
|
#endif
|
|
break;
|
|
}
|
|
if (error == ERROR_ACCESS_DENIED) {
|
|
std::cerr << "[ERROR] Cannot query the memory region. " << std::hex << start_va << " Error: " << std::dec << error << std::endl;
|
|
break;
|
|
}
|
|
if (!is_read) {
|
|
// on any other error:
|
|
std::cerr << "[WARNING] Cannot query the memory region. " << std::hex<< start_va << " Error: " << std::dec << error << ", skipping the page..." << std::endl;
|
|
start_va += PAGE_SIZE;
|
|
continue;
|
|
}
|
|
if ((page_info.State & MEM_FREE) || (page_info.State & MEM_COMMIT) == 0) {
|
|
if (page_info.RegionSize != 0) {
|
|
//std::cout << "Free: " << std::hex << start_va << " RegionSize:" << page_info.RegionSize << std::endl;
|
|
start_va += page_info.RegionSize;
|
|
continue;
|
|
}
|
|
}
|
|
if (page_info.RegionSize == 0) {
|
|
start_va += PAGE_SIZE;
|
|
continue;
|
|
}
|
|
//std::cout << "Commited: " << std::hex << start_va << " RegionSize: " << page_info.RegionSize << " Err: " << std::dec << error << std::endl;
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
};
|
|
};
|
|
|
|
size_t pesieve::util::enum_workingset(HANDLE processHandle, std::set<mem_region_info> ®ion_bases)
|
|
{
|
|
region_bases.clear();
|
|
|
|
MEMORY_BASIC_INFORMATION page_info = { 0 };
|
|
ULONGLONG next_va = 0;
|
|
while (get_next_commited_region(processHandle, next_va, page_info))
|
|
{
|
|
ULONGLONG base = (ULONGLONG)page_info.BaseAddress;
|
|
next_va = base + page_info.RegionSize; //end of the region
|
|
mem_region_info curr_info((ULONGLONG)page_info.AllocationBase, base, page_info.RegionSize);
|
|
|
|
if (region_bases.find(curr_info) != region_bases.end()) {
|
|
// don't let it stuck on adding the same region over and over again
|
|
break;
|
|
}
|
|
region_bases.insert(curr_info);
|
|
}
|
|
return region_bases.size();
|
|
}
|
|
|
|
DWORD pesieve::util::count_workingset_entries(HANDLE processHandle)
|
|
{
|
|
DWORD number_of_entries = 1;
|
|
DWORD buffer_size = sizeof(PSAPI_WORKING_SET_INFORMATION) + (number_of_entries * sizeof(PSAPI_WORKING_SET_BLOCK));
|
|
PSAPI_WORKING_SET_INFORMATION* buffer = reinterpret_cast<PSAPI_WORKING_SET_INFORMATION*>(calloc(1, buffer_size));
|
|
if (!buffer) {
|
|
return 0; //this should not happen
|
|
}
|
|
DWORD res = QueryWorkingSet(processHandle, buffer, buffer_size);
|
|
if (res == FALSE && GetLastError() == ERROR_BAD_LENGTH) {
|
|
// ERROR_BAD_LENGTH is normal: we didn't provide the buffer that could fit all the entries
|
|
res = TRUE;
|
|
}
|
|
number_of_entries = static_cast<DWORD>(buffer->NumberOfEntries);
|
|
free(buffer); buffer = NULL;
|
|
|
|
if (!res) {
|
|
return 0;
|
|
}
|
|
#ifdef _DEBUG
|
|
std::cout << "Number of entries: " << std::dec << number_of_entries << std::endl;
|
|
#endif
|
|
return number_of_entries;
|
|
}
|