WIP Patch

This commit is contained in:
Guus Waals 2025-06-06 20:09:23 +08:00
parent 72c6656e01
commit 8b73364db9
1 changed files with 136 additions and 24 deletions

View File

@ -117,7 +117,8 @@ struct Patcher {
ptr--;
}
return availableSpace;
// Keep 1 byte marging, might be another string or terminator
return availableSpace - 1;
}
void logMainFunctionCode() {
@ -181,16 +182,111 @@ struct Patcher {
return true;
}
// Tries to find an import in the input PE file
void findImportInPE(std::string_view name) {
auto &imports = peReader.get_imports();
for (auto &import : imports) {
if (import.get_name() == name) {
spdlog::info("Found import: {}", name);
return;
// Tries to find an import in the input PE file and return its IAT address
uint32_t findImportInPE(const std::string &functionName) {
auto &directories = peReader.get_directories();
// Import table is directory entry 1
if (directories.size() <= 1 || !directories[1]) {
spdlog::error("No import directory found in PE file");
return 0;
}
uint32_t importRVA = directories[1]->get_virtual_address();
uint32_t importSize = directories[1]->get_size();
if (importRVA == 0 || importSize == 0) {
spdlog::error("Invalid import directory");
return 0;
}
spdlog::info("Import directory at RVA: 0x{:x}, size: 0x{:x}", importRVA,
importSize);
// Find which section contains the import table
COFFI::section *importSection = nullptr;
uint32_t importFileOffset = 0;
auto &sections = peReader.get_sections();
for (auto &section : sections) {
uint32_t sectionRVA = section->get_virtual_address();
uint32_t sectionSize = section->get_virtual_size();
if (importRVA >= sectionRVA && importRVA < sectionRVA + sectionSize) {
importSection = section;
importFileOffset = importRVA - sectionRVA;
break;
}
}
spdlog::error("Could not find import: {}", name);
if (!importSection) {
spdlog::error("Could not find section containing import table");
return 0;
}
// Parse import descriptors
const char *sectionData = importSection->get_data();
const char *importData = sectionData + importFileOffset;
// Import descriptor structure (20 bytes each)
struct ImportDescriptor {
uint32_t originalFirstThunk; // RVA to import lookup table
uint32_t timeDateStamp;
uint32_t forwarderChain;
uint32_t nameRVA; // RVA to DLL name
uint32_t firstThunk; // RVA to import address table
};
const ImportDescriptor *desc =
reinterpret_cast<const ImportDescriptor *>(importData);
// Iterate through import descriptors
for (int i = 0; desc[i].nameRVA != 0; i++) {
// Get DLL name
std::string dllName = getRVAString(desc[i].nameRVA);
spdlog::info("Checking imports from DLL: {}", dllName);
// Parse the import lookup table
uint32_t lookupRVA = desc[i].originalFirstThunk;
uint32_t iatRVA = desc[i].firstThunk;
if (lookupRVA == 0)
lookupRVA = iatRVA; // Use IAT if no lookup table
const uint32_t *lookupTable =
reinterpret_cast<const uint32_t *>(getRVAData(lookupRVA));
const uint32_t *iatTable =
reinterpret_cast<const uint32_t *>(getRVAData(iatRVA));
if (!lookupTable || !iatTable)
continue;
// Iterate through function imports
for (int j = 0; lookupTable[j] != 0; j++) {
uint32_t nameRVA = lookupTable[j];
// Skip ordinal imports (high bit set)
if (nameRVA & 0x80000000)
continue;
// Get function name (skip hint, first 2 bytes)
const char *funcNamePtr = getRVAString(nameRVA + 2);
if (!funcNamePtr)
continue;
std::string funcName(funcNamePtr);
if (funcName == functionName) {
uint32_t iatAddress = imageBase + iatRVA + (j * sizeof(uint32_t));
spdlog::info("Found import '{}' in {} at IAT address: 0x{:x}",
functionName, dllName, iatAddress);
return iatAddress;
}
}
}
spdlog::error("Could not find import: {}", functionName);
return 0;
}
bool processRelocations(char *textStart, char *textEnd, char *&rdataPtr) {
@ -223,18 +319,15 @@ struct Patcher {
spdlog::info("Processing relocation at VA 0x{:x} for symbol: {}",
relocRVA + imageBase, symbolName);
if (symbolName.starts_with("__imp_")) {
if (symbolName.starts_with("__imp__")) {
// Import symbol handling
std::string functionName = symbolName.substr(6);
std::string functionName = symbolName.substr(7);
size_t atPos = functionName.find('@');
if (atPos != std::string::npos) {
functionName = functionName.substr(0, atPos);
}
spdlog::info("Looking for import function: {}", functionName);
spdlog::warn("Import resolution not fully implemented yet for: {}",
functionName);
resolvedAddress = 0x12345678; // Placeholder
resolvedAddress = findImportInPE(functionName);
} else if (symbolName.starts_with("??_C@")) {
// String constant handling
@ -252,10 +345,7 @@ struct Patcher {
uint32_t constOffset = constSymbol->get_value();
// Find string length
uint32_t stringLen = 0;
while (constData[constOffset + stringLen] != 0) {
stringLen++;
}
uint32_t stringLen = strlen(constData + constOffset);
stringLen++; // Include null terminator
if (rdataPtr + stringLen > rdataEndPtr) {
@ -287,7 +377,7 @@ struct Patcher {
ptr[0] = resolvedAddress;
spdlog::info(
"Applied relocation: patched offset 0x{:x} with address 0x{:x}",
patchOffset, resolvedAddress);
relocRVA + imageBase, resolvedAddress);
}
}
@ -339,16 +429,16 @@ struct Patcher {
reinterpret_cast<const uint8_t *>(mainSection->get_data()) + mainOffset;
std::memcpy(dstText, mainCode, mainSize);
spdlog::info("Injected code into existing .text section space (size "
"unchanged: 0x{:x})",
textSize);
// Process relocations
char *rdataPtr = dstRdata;
if (!processRelocations(dstText, dstText + mainSize, rdataPtr)) {
return false;
}
spdlog::info("Injected code into existing .text section space (size "
"unchanged: 0x{:x})",
textSize);
// Save the modified PE file
spdlog::info("Saving patched PE file to: {}", outputFile);
if (!peReader.save(outputFile)) {
@ -379,6 +469,28 @@ struct Patcher {
return true;
}
// Helper function to get string at RVA
const char *getRVAString(uint32_t rva) {
return reinterpret_cast<const char *>(getRVAData(rva));
}
// Helper function to get data pointer at RVA
const void *getRVAData(uint32_t rva) {
auto &sections = peReader.get_sections();
for (auto &section : sections) {
uint32_t sectionRVA = section->get_virtual_address();
uint32_t sectionSize = section->get_virtual_size();
if (rva >= sectionRVA && rva < sectionRVA + sectionSize) {
uint32_t offset = rva - sectionRVA;
return section->get_data() + offset;
}
}
return nullptr;
}
};
int main(int argc, char *argv[]) {