reman3/patcher/patcher.cpp

396 lines
12 KiB
C++

#include <windows.h>
#include <coffi/coffi.hpp>
#include <CLI11.hpp>
#include <spdlog/spdlog.h>
struct Patcher {
COFFI::coffi objReader;
COFFI::coffi peReader;
std::string inputFile;
std::string outputFile;
COFFI::symbol *mainSymbol = nullptr;
COFFI::section *mainSection = nullptr;
COFFI::section *textSection = nullptr;
COFFI::section *rdataSection = nullptr;
uint64_t imageBase = 0;
uint32_t mainOffset = 0;
uint32_t mainSize = 0;
bool loadFiles() {
std::string objPath = SRC_OBJECT;
SPDLOG_INFO("Loading object file: {}", objPath);
if (!objReader.load(objPath)) {
spdlog::error("Failed to load object file: {}", objPath);
return false;
}
SPDLOG_INFO("Loading PE file: {}", inputFile);
if (!peReader.load(inputFile)) {
spdlog::error("Failed to load PE file: {}", inputFile);
return false;
}
// Get image base
auto winHeader = peReader.get_win_header();
if (winHeader) {
imageBase = winHeader->get_image_base();
}
spdlog::info("PE Image base: 0x{:x}", imageBase);
return true;
}
bool findMainSymbol() {
auto &symbols = *objReader.get_symbols();
for (auto &sym : symbols) {
SPDLOG_INFO("Symbol: {}", sym.get_name());
if (sym.get_name() == "_ref") {
mainSymbol = &sym;
break;
}
}
if (!mainSymbol) {
spdlog::error("Could not find 'main' symbol in object file");
return false;
}
// Get the section containing the main function
auto &sections = objReader.get_sections();
mainSection = sections[mainSymbol->get_section_number() - 1];
mainOffset = mainSymbol->get_value();
return true;
}
uint32_t calculateMainSize() {
auto &symbols = *objReader.get_symbols();
uint32_t nextSymbolOffset = UINT32_MAX;
// Find the next symbol in the same section to calculate size
for (auto &sym : symbols) {
if (sym.get_section_number() == mainSymbol->get_section_number() &&
sym.get_value() > mainOffset && sym.get_value() < nextSymbolOffset) {
nextSymbolOffset = sym.get_value();
}
}
if (nextSymbolOffset != UINT32_MAX) {
mainSize = nextSymbolOffset - mainOffset;
spdlog::info(
"Calculated main function size: {} bytes (next symbol at offset {})",
mainSize, nextSymbolOffset);
} else {
// If no next symbol found, use remaining section size
mainSize = mainSection->get_data_size() - mainOffset;
spdlog::info(
"No next symbol found, using remaining section size: {} bytes",
mainSize);
}
return mainSize;
}
COFFI::section *findSection(const std::string &name, bool isPE = true) {
auto &sections = isPE ? peReader.get_sections() : objReader.get_sections();
for (auto &section : sections) {
if (section->get_name() == name) {
return section;
}
}
return nullptr;
}
uint32_t findAvailableSpaceAtEnd(COFFI::section *section) {
auto sectionData = section->get_data();
uint32_t sectionSize = section->get_data_size();
const char *sectionEnd = sectionData + sectionSize;
uint32_t availableSpace = 0;
// Search backwards from the end to find contiguous null bytes
const char *ptr = sectionEnd-1;
while (ptr > sectionData) {
if (*ptr != 0x00) {
break;
}
availableSpace++;
ptr--;
}
// }
// for (int32_t i = sectionSize - 1; i >= 0; i--) {
// if (reinterpret_cast<const uint8_t*>(sectionData)[i] == 0x00) {
// availableSpace++;
// } else {
// break; // Found non-null byte, stop counting
// }
// }
return availableSpace;
}
void logMainFunctionCode() {
auto mainCodeData = mainSection->get_data();
spdlog::info("Found main function at offset {} with size {}", mainOffset,
mainSize);
spdlog::info("Main function code:");
std::string s;
for (uint32_t i = 0; i < mainSize; i++) {
if (i > 0 && i % 16 == 0) {
spdlog::info("{}", s);
s.clear();
}
if (s.size() > 0)
s += " ";
s += fmt::format("{:02X}", mainCodeData[i]);
}
if (s.size() > 0)
spdlog::info("{}", s);
}
void logRelocations() {
auto &sectionRelocations = mainSection->get_relocations();
spdlog::info("Relocations for section {} (containing main function):",
mainSymbol->get_section_number());
for (auto &reloc : sectionRelocations) {
spdlog::info(" Relocation at offset 0x{:x}, type: {}, symbol index: {}",
reloc.get_virtual_address(), reloc.get_type(),
reloc.get_symbol_table_index());
spdlog::info(" -> Symbol: {}", reloc.get_symbol());
}
}
void logSectionData(COFFI::section *section) {
uint32_t sectionRVA = section->get_virtual_address();
uint32_t sectionSize = section->get_virtual_size();
uint64_t sectionVA = imageBase + sectionRVA;
uint64_t sectionEndVA = sectionVA + sectionSize;
spdlog::info("Section {} RVA: 0x{:x}, size: 0x{:x}", section->get_name(),
sectionRVA, sectionSize);
spdlog::info("Section {} VA: 0x{:x} - 0x{:x}", section->get_name(),
sectionVA, sectionEndVA);
}
bool validateSections() {
textSection = findSection(".text");
if (!textSection) {
spdlog::error("Could not find .text section in PE file");
return false;
}
logSectionData(textSection);
rdataSection = findSection(".rdata");
if (!rdataSection) {
spdlog::error("Could not find .rdata section in PE file");
return false;
}
logSectionData(rdataSection);
return true;
}
bool processRelocations(std::vector<uint8_t> &newTextData,
std::vector<uint8_t> &newRdataData,
uint32_t injectionOffset,
uint32_t &rdataCurrentOffset) {
auto &symbols = *objReader.get_symbols();
auto &sections = objReader.get_sections();
auto &sectionRelocations = mainSection->get_relocations();
uint32_t rdataSize = rdataSection->get_data_size();
uint32_t rdataRVA = rdataSection->get_virtual_address();
spdlog::info("Processing relocations...");
for (auto &reloc : sectionRelocations) {
if (reloc.get_type() != 6)
continue; // Only handle type 6 (IMAGE_REL_I386_DIR32)
uint32_t relocOffset = reloc.get_virtual_address();
std::string symbolName = reloc.get_symbol();
uint32_t resolvedAddress = 0;
spdlog::info("Processing relocation at offset 0x{:x} for symbol: {}",
relocOffset, symbolName);
if (symbolName.starts_with("__imp_")) {
// Import symbol handling
std::string functionName = symbolName.substr(6);
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
} else if (symbolName.starts_with("??_C@")) {
// String constant handling
COFFI::symbol *constSymbol = nullptr;
for (auto &sym : symbols) {
if (sym.get_name() == symbolName) {
constSymbol = &sym;
break;
}
}
if (constSymbol) {
auto constSection = sections[constSymbol->get_section_number() - 1];
auto constData = constSection->get_data();
uint32_t constOffset = constSymbol->get_value();
// Find string length
uint32_t stringLen = 0;
while (constData[constOffset + stringLen] != 0) {
stringLen++;
}
stringLen++; // Include null terminator
if (rdataCurrentOffset + stringLen > rdataSize) {
spdlog::error("Not enough space in .rdata for string constant");
return false;
}
// Copy string to .rdata
std::memcpy(newRdataData.data() + rdataCurrentOffset,
constData + constOffset, stringLen);
resolvedAddress = imageBase + rdataRVA + rdataCurrentOffset;
spdlog::info("Copied string constant '{}' to .rdata at RVA 0x{:x} "
"(VA: 0x{:x})",
std::string(constData + constOffset),
rdataRVA + rdataCurrentOffset, resolvedAddress);
rdataCurrentOffset += stringLen;
}
}
// Apply the relocation
if (resolvedAddress != 0) {
uint32_t patchOffset = injectionOffset + relocOffset;
*reinterpret_cast<uint32_t *>(newTextData.data() + patchOffset) =
resolvedAddress;
spdlog::info(
"Applied relocation: patched offset 0x{:x} with address 0x{:x}",
patchOffset, resolvedAddress);
}
}
return true;
}
bool patchAndSave() {
// Check available space in both sections
uint32_t textAvailableSpace = findAvailableSpaceAtEnd(textSection);
uint32_t rdataAvailableSpace = findAvailableSpaceAtEnd(rdataSection);
spdlog::info("Found {} bytes of available space (null bytes) at end of "
".text section",
textAvailableSpace);
spdlog::info("Found {} bytes of available space in .rdata section",
rdataAvailableSpace);
if (textAvailableSpace < mainSize) {
spdlog::error(
"Not enough space in .text section! Need {} bytes, found {} bytes",
mainSize, textAvailableSpace);
return false;
}
// Calculate injection points
uint32_t textSize = textSection->get_data_size();
uint32_t rdataSize = rdataSection->get_data_size();
uint32_t injectionOffset = textSize - textAvailableSpace;
uint32_t rdataInjectionOffset = rdataSize - rdataAvailableSpace;
uint32_t rdataCurrentOffset = rdataInjectionOffset;
uint64_t textSectionVA = imageBase + textSection->get_virtual_address();
uint64_t injectionVA = textSectionVA + injectionOffset;
spdlog::info(
"Injecting {} bytes at .text section offset 0x{:x} (VA: 0x{:x})",
mainSize, injectionOffset, injectionVA);
// Create copies of section data
auto textSectionData = textSection->get_data();
auto rdataData = rdataSection->get_data();
std::vector<uint8_t> newTextData(
reinterpret_cast<const uint8_t *>(textSectionData),
reinterpret_cast<const uint8_t *>(textSectionData) + textSize);
std::vector<uint8_t> newRdataData(
reinterpret_cast<const uint8_t *>(rdataData),
reinterpret_cast<const uint8_t *>(rdataData) + rdataSize);
// Copy main function code
const uint8_t *mainCode =
reinterpret_cast<const uint8_t *>(mainSection->get_data()) + mainOffset;
std::memcpy(newTextData.data() + injectionOffset, mainCode, mainSize);
// Process relocations
if (!processRelocations(newTextData, newRdataData, injectionOffset,
rdataCurrentOffset)) {
return false;
}
// Update sections
rdataSection->set_data(reinterpret_cast<const char *>(newRdataData.data()),
newRdataData.size());
textSection->set_data(reinterpret_cast<const char *>(newTextData.data()),
newTextData.size());
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)) {
spdlog::error("Failed to save patched PE file to: {}", outputFile);
return false;
}
spdlog::info(
"Successfully patched PE file! Main function injected at VA: 0x{:x}",
injectionVA);
return true;
}
bool run() {
if (!loadFiles())
return false;
if (!findMainSymbol())
return false;
calculateMainSize();
logMainFunctionCode();
logRelocations();
if (!validateSections())
return false;
if (!patchAndSave())
return false;
return true;
}
};
int main(int argc, char *argv[]) {
CLI::App app("Patcher");
std::string inputFile;
std::string outputFile;
app.add_option("-i,--input", inputFile, "Input exe file to patch")
->required();
app.add_option("-o,--output", outputFile, "Output patched exe file")
->required();
CLI11_PARSE(app, argc, argv);
Patcher patcher;
patcher.inputFile = inputFile;
patcher.outputFile = outputFile;
return patcher.run() ? 0 : 1;
}