reman3/patcher/patcher.cpp

#include <windows.h>
#include <coffi/coffi.hpp>
#include <CLI11.hpp>
#include <spdlog/spdlog.h>

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);

  // Load the object file containing the main function
  COFFI::coffi objReader;
  std::string objPath = SRC_OBJECT;
  SPDLOG_INFO("Loading object file: {}", objPath);
  if (!objReader.load(objPath)) {
    spdlog::error("Failed to load object file: {}", objPath);
    return 1;
  }

  // Load the source PE file
  COFFI::coffi peReader;
  SPDLOG_INFO("Loading PE file: {}", inputFile);
  if (!peReader.load(inputFile)) {
    spdlog::error("Failed to load PE file: {}", inputFile);
    return 1;
  }

  // Find the 'main' function in the object file
  auto &symbols = *objReader.get_symbols();
  COFFI::symbol *mainSymbol = nullptr;
  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 1;
  }

  // Get the section containing the main function
  auto &sections = objReader.get_sections();
  auto mainSection = sections[mainSymbol->get_section_number() - 1];

  // Calculate main function size using next symbol method
  uint32_t mainOffset = mainSymbol->get_value();
  uint32_t mainSize = 0;

  // Find the next symbol in the same section to calculate size
  uint32_t nextSymbolOffset = UINT32_MAX;
  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);
  }

  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);

  // Find .text section in PE file
  auto &peSections = peReader.get_sections();
  COFFI::section *textSection = nullptr;
  for (auto &section : peSections) {
    if (section->get_name() == ".text") {
      textSection = section;
      break;
    }
  }

  if (!textSection) {
    spdlog::error("Could not find .text section in PE file");
    return 1;
  }

  uint32_t textSectionEnd =
      textSection->get_virtual_address() + textSection->get_virtual_size();
  spdlog::info("Found .text section, end at virtual address: 0x{:x}",
               textSectionEnd);

  return 0;
}