gbe_fork/helpers/pe_helpers.cpp

473 lines
14 KiB
C++
Raw Normal View History

2024-01-12 21:37:22 +08:00
#include "pe_helpers/pe_helpers.hpp"
#include "common_helpers/common_helpers.hpp"
2024-01-12 21:37:22 +08:00
#include <vector>
#include <utility>
#include <mutex>
#include <cwchar>
2024-01-12 21:37:22 +08:00
static inline bool is_hex(const char c)
{
return (c >= '0' && c <= '9') ||
(c >= 'a' && c <= 'f') ||
(c >= 'A' && c <= 'F');
}
static inline uint8_t char_to_byte(const char c)
{
if (c >= '0' && c <= '9') return (uint8_t)(c - '0');
if (c >= 'a' && c <= 'f') return (uint8_t)(c - 'a') + 0xa;
if (c >= 'A' && c <= 'F') return (uint8_t)(c - 'A') + 0xa;
return (uint8_t)c;
}
PIMAGE_NT_HEADERS pe_helpers::get_nt_header(HMODULE hModule)
2024-01-12 21:37:22 +08:00
{
// https://dev.to/wireless90/validating-the-pe-signature-my-av-flagged-me-windows-pe-internals-2m5o/
PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)(char*)hModule;
if (dosHeader->e_magic != 0x5A4D) { // "MZ"
return nullptr;
}
2024-01-12 21:37:22 +08:00
LONG newExeHeaderOffset = dosHeader->e_lfanew;
return (PIMAGE_NT_HEADERS)((char*)hModule + newExeHeaderOffset);
}
PIMAGE_FILE_HEADER pe_helpers::get_file_header(HMODULE hModule)
2024-01-12 21:37:22 +08:00
{
auto nt_header = get_nt_header(hModule);
return nt_header ? &nt_header->FileHeader : nullptr;
2024-01-12 21:37:22 +08:00
}
PIMAGE_OPTIONAL_HEADER pe_helpers::get_optional_header(HMODULE hModule)
2024-01-12 21:37:22 +08:00
{
auto nt_header = get_nt_header(hModule);
return nt_header ? &nt_header->OptionalHeader : nullptr;
2024-01-12 21:37:22 +08:00
}
uint8_t* pe_helpers::search_memory(uint8_t *mem, size_t size, const std::string &search_patt)
{
if (!mem || !size || search_patt.empty()) return nullptr;
if (search_patt.find_first_not_of(" \t", 0) == std::string::npos) {
return nullptr; // empty search patt
}
const uint8_t *end = mem + size;
for (uint8_t *base = mem; base < end; ++base)
{
// incremental offset after each byte match
size_t search_patt_offset = 0;
bool error = false;
for (const uint8_t *displacement = base; displacement < end; ++displacement)
{
uint8_t mask = 0xFF;
uint8_t s_byte = 0;
// skip spaces
search_patt_offset = search_patt.find_first_not_of(" \t", search_patt_offset);
if (search_patt_offset == std::string::npos) {
break;
}
const auto this_char = search_patt[search_patt_offset];
const auto next_char = (search_patt_offset + 1) < search_patt.size()
? search_patt[search_patt_offset + 1]
: '\0';
if (this_char == '?') {
if (next_char == '?' || // "??"
next_char == ' ' || // "? "
next_char == 't') { // "? "
mask = 0x00;
s_byte = 0;
} else if (is_hex(next_char)) { // "?c"
mask = 0x0F;
s_byte = char_to_byte(next_char);
} else { // unknown
return nullptr;
}
// skip
search_patt_offset += 2;
} else if (is_hex(this_char)) {
if (next_char == '?') { // "c?"
mask = 0xF0;
s_byte = char_to_byte(this_char) << 4;
} else if (is_hex(next_char)) { // "34"
mask = 0xFF;
s_byte = (char_to_byte(this_char) << 4) | char_to_byte(next_char);
} else { // unknown
return nullptr;
}
// skip
search_patt_offset += 2;
} else { // unknown
return nullptr;
}
if ((*displacement & mask) != (s_byte & mask)) {
error = true;
break;
}
}
if (!error && (search_patt_offset >= search_patt.size())) {
return base;
}
}
return nullptr;
}
bool pe_helpers::replace_memory(uint8_t *mem, size_t size, const std::string &replace_patt, HANDLE hProcess)
{
if (!mem || !size || replace_patt.empty()) return false;
size_t replace_patt_offset = replace_patt.find_first_not_of(" \t", 0);
if (replace_patt_offset == std::string::npos) {
return false; // empty patt
}
// mask - byte
std::vector<std::pair<uint8_t, uint8_t>> replace_bytes{};
for (;
replace_patt_offset < replace_patt.size();
replace_patt_offset = replace_patt.find_first_not_of(" \t", replace_patt_offset)) {
const auto this_char = replace_patt[replace_patt_offset];
const auto next_char = (replace_patt_offset + 1) < replace_patt.size()
? replace_patt[replace_patt_offset + 1]
: '\0';
if (this_char == '?') {
if (next_char == '?' || // "??"
next_char == ' ' || // "? "
next_char == 't') { // "? "
replace_bytes.push_back({
0x00,
0,
});
} else if (is_hex(next_char)) { // "?c"
replace_bytes.push_back({
0x0F,
char_to_byte(next_char),
});
} else { // unknown
return false;
}
// skip
replace_patt_offset += 2;
} else if (is_hex(this_char)) {
if (next_char == '?') { // "c?"
replace_bytes.push_back({
0xF0,
char_to_byte(this_char) << 4,
});
} else if (is_hex(next_char)) { // "34"
replace_bytes.push_back({
0xFF,
(char_to_byte(this_char) << 4) | char_to_byte(next_char),
});
} else { // unknown
return false;
}
// skip
replace_patt_offset += 2;
} else { // unknown
return false;
}
}
// remove trailing "??"
// while last element == "??"
while (replace_bytes.size() &&
replace_bytes.back().first == 0x00)
{
replace_bytes.pop_back();
}
if (replace_bytes.empty() || replace_bytes.size() > size) return false;
// change protection
DWORD current_protection = 0;
if (!VirtualProtectEx(hProcess, mem, replace_bytes.size(), PAGE_READWRITE, &current_protection)) {
return false;
}
for (auto &rp : replace_bytes) {
if (rp.first == 0x00) {
++mem;
continue;
}
2024-01-12 21:37:22 +08:00
const uint8_t b_mem = (uint8_t)(*mem & (uint8_t)~rp.first);
const uint8_t b_replace = (uint8_t)(rp.second & rp.first);
const uint8_t new_b_mem = b_mem | b_replace;
*mem = b_mem | b_replace;
++mem;
}
// restore protection
if (!VirtualProtectEx(hProcess, mem, replace_bytes.size(), current_protection, &current_protection)) {
return false;
}
return true;
}
// https://learn.microsoft.com/en-us/windows/win32/debug/retrieving-the-last-error-code
std::string pe_helpers::get_err_string(DWORD code)
{
std::string err_str(8192, '\0');
DWORD msg_chars = FormatMessageA(
FORMAT_MESSAGE_FROM_SYSTEM |
FORMAT_MESSAGE_IGNORE_INSERTS,
NULL,
code,
MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL),
(LPSTR)&err_str[0],
err_str.size(),
NULL);
if (!msg_chars) return std::string();
err_str = err_str.substr(0, msg_chars);
return err_str;
}
bool pe_helpers::is_module_64(HMODULE hModule)
{
auto file_header = get_file_header(hModule);
return file_header ? (file_header->Machine == IMAGE_FILE_MACHINE_AMD64) : false;
2024-01-12 21:37:22 +08:00
}
bool pe_helpers::is_module_32(HMODULE hModule)
{
auto file_header = get_file_header(hModule);
return file_header ? (file_header->Machine == IMAGE_FILE_MACHINE_I386) : false;
2024-01-12 21:37:22 +08:00
}
pe_helpers::SectionHeadersResult pe_helpers::get_section_headers(HMODULE hModule)
{
PIMAGE_NT_HEADERS ntHeader = get_nt_header(hModule);
PIMAGE_OPTIONAL_HEADER optionalHeader = &ntHeader->OptionalHeader;
PIMAGE_FILE_HEADER fileHeader = get_file_header(hModule);
WORD optionalHeadrSize = fileHeader->SizeOfOptionalHeader;
struct SectionHeadersResult res {};
res.count = fileHeader->NumberOfSections;
res.ptr = fileHeader->NumberOfSections
? (PIMAGE_SECTION_HEADER)((char *)optionalHeader + optionalHeadrSize)
: nullptr;
return res;
}
PIMAGE_SECTION_HEADER pe_helpers::get_section_header_with_name(HMODULE hModule, const char *name)
{
if (!name) return nullptr;
auto res = get_section_headers(hModule);
if (!res.count) return nullptr;
for (size_t i = 0; i < res.count; ++i) {
if (strncmp((const char *)res.ptr[i].Name, name, sizeof(res.ptr[i].Name)) == 0) {
return &res.ptr[i];
}
}
return nullptr;
}
2024-01-25 05:11:11 +08:00
DWORD pe_helpers::loadlib_remote(HANDLE hProcess, const std::wstring &lib_fullpath, const char** err_reason)
{
// create a remote page
const size_t lib_path_str_bytes = lib_fullpath.size() * sizeof(lib_fullpath[0]);
LPVOID lib_remote_page = VirtualAllocEx(
2024-01-12 21:37:22 +08:00
hProcess,
NULL,
lib_path_str_bytes + sizeof(lib_fullpath[0]) * 2, // *2 just to be safe
MEM_RESERVE | MEM_COMMIT,
PAGE_READWRITE
2024-01-25 05:11:11 +08:00
);
2024-01-12 21:37:22 +08:00
2024-01-25 05:11:11 +08:00
if (!lib_remote_page) {
if (err_reason) {
*err_reason = "Failed to remotely allocate page with VirtualAllocEx()";
}
return GetLastError();
2024-01-12 21:37:22 +08:00
}
2024-01-25 05:11:11 +08:00
SIZE_T bytes_written = 0;
BOOL written = WriteProcessMemory(
hProcess,
lib_remote_page,
(LPCVOID)&lib_fullpath[0],
lib_path_str_bytes,
&bytes_written
);
if (!written || bytes_written < lib_path_str_bytes) {
// cleanup allcoated page
VirtualFreeEx(
hProcess,
lib_remote_page,
0,
MEM_RELEASE);
if (err_reason) {
*err_reason = "Failed to remotely write dll path with WriteProcessMemory()";
}
return GetLastError();
}
2024-01-12 21:37:22 +08:00
2024-01-25 05:11:11 +08:00
// call LoadLibraryW() and pass the dll fullpath
HANDLE remote_thread = CreateRemoteThread(
hProcess,
NULL,
0,
(LPTHREAD_START_ROUTINE)LoadLibraryW,
lib_remote_page,
0,
NULL);
if (!remote_thread) {
// cleanup allcoated page
VirtualFreeEx(
hProcess,
lib_remote_page,
0,
MEM_RELEASE);
2024-01-12 21:37:22 +08:00
2024-01-25 05:11:11 +08:00
if (err_reason) {
*err_reason = "Failed to create/run remote thread with CreateRemoteThread()";
}
return GetLastError();
2024-01-12 21:37:22 +08:00
}
2024-01-25 05:11:11 +08:00
// wait for DllMain
WaitForSingleObject(remote_thread, INFINITE);
CloseHandle(remote_thread);
2024-01-12 21:37:22 +08:00
// cleanup allcoated page
VirtualFreeEx(
2024-01-25 05:11:11 +08:00
hProcess,
lib_remote_page,
0,
MEM_RELEASE);
2024-01-12 21:37:22 +08:00
2024-01-25 05:11:11 +08:00
return ERROR_SUCCESS;
2024-01-12 21:37:22 +08:00
}
size_t pe_helpers::get_pe_size(HMODULE hModule)
{
// https://stackoverflow.com/a/34695773
// https://learn.microsoft.com/en-us/windows/win32/debug/pe-format
// "The PE file header consists of a Microsoft MS-DOS stub, the PE signature, the COFF file header, and an optional header"
// "The combined size of an MS-DOS stub, PE header, and section headers rounded up to a multiple of FileAlignment."
size_t size = get_optional_header(hModule)->SizeOfHeaders;
SectionHeadersResult headers = get_section_headers(hModule);
for (size_t i = 0; i < headers.count; ++i) {
size += headers.ptr[i].SizeOfRawData;
}
return size;
}
static std::wstring path_w{};
static std::string path_a{};
const std::string& pe_helpers::get_current_exe_path()
{
if (path_a.empty()) {
get_current_exe_path_w();
}
return path_a;
}
const std::wstring& pe_helpers::get_current_exe_path_w()
{
static std::recursive_mutex path_mtx{};
if (path_w.empty()) {
std::lock_guard lk(path_mtx);
if (path_w.empty()) {
DWORD err = GetLastError();
path_w.resize(8192);
DWORD read_chars = GetModuleFileNameW(GetModuleHandleW(nullptr), &path_w[0], (DWORD)path_w.size());
if (read_chars >= path_w.size()) {
path_w.resize(read_chars);
read_chars = GetModuleFileNameW(GetModuleHandleW(nullptr), &path_w[0], (DWORD)path_w.size());
}
if ((read_chars < path_w.size()) && path_w[0]) {
path_w = path_w.substr(0, path_w.find_last_of(L"\\/") + 1);
auto cvt_state = std::mbstate_t();
const wchar_t* src = &path_w[0];
size_t conversion_bytes = std::wcsrtombs(nullptr, &src, 0, &cvt_state);
path_a.resize(conversion_bytes + 1);
std::wcsrtombs(&path_a[0], &src, path_a.size(), &cvt_state);
path_a = path_a.substr(0, conversion_bytes);
} else {
path_w.clear();
}
SetLastError(err);
}
}
return path_w;
}
bool pe_helpers::ends_with_i(PUNICODE_STRING target, const std::wstring &query)
{
return common_helpers::ends_with_i(
std::wstring(target->Buffer, (PWSTR)((char*)target->Buffer + target->Length)),
query
);
}
MEMORY_BASIC_INFORMATION pe_helpers::get_mem_page_details(const void* mem)
{
MEMORY_BASIC_INFORMATION mbi{};
if (VirtualQuery(mem, &mbi, sizeof(mbi))) {
return mbi;
} else {
return {};
}
}
size_t pe_helpers::get_current_exe_mem_size()
{
auto hmod = GetModuleHandleW(NULL);
size_t size = 0;
{
MEMORY_BASIC_INFORMATION mbi{};
if (!VirtualQuery((LPVOID)hmod, &mbi, sizeof(mbi))) {
return 0;
}
size = mbi.RegionSize; // PE header
}
auto sections = get_section_headers(hmod);
if (!sections.count) {
return 0;
}
for (size_t i = 0; i < sections.count; ++i) {
auto section = sections.ptr[i];
MEMORY_BASIC_INFORMATION mbi{};
if (!VirtualQuery((LPCVOID)((uint8_t *)hmod + section.VirtualAddress), &mbi, sizeof(mbi))) {
return 0;
}
size = mbi.RegionSize; // actual section size in mem
}
return size;
}