reman3/patcher/patcher.cpp

#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--;
    }

    // Keep 1 byte margin, might be another string or terminator
    return availableSpace > 0 ? availableSpace - 1 : 0;
  }

  void logMainFunctionCode(COFFI::section *objMainSection, uint32_t functionOffset, uint32_t functionSize) {
    auto mainCodeData = objMainSection->get_data();
    spdlog::info("Found main function at offset {} with size {}", functionOffset, functionSize);
    spdlog::info("Main function code:");

    std::string s;
    for (uint32_t i = 0; i < functionSize; i++) {
      if (i > 0 && i % 16 == 0) {
        spdlog::info("{}", s);
        s.clear();
      }
      if (s.size() > 0)
        s += " ";
      s += fmt::format("{:02X}", mainCodeData[functionOffset + i]);
    }
    if (s.size() > 0)
      spdlog::info("{}", s);
  }

  void logRelocations(COFFI::section *objMainSection, uint32_t symbolSectionNumber) {
    auto &sectionRelocations = objMainSection->get_relocations();
    spdlog::info("Relocations for section {} (containing main function):", symbolSectionNumber);
    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;
  }

  // Returns the VA of the import address table entry for the specified function
  uint64_t findImportVA(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 importTableRVA = directories[1]->get_virtual_address();
    uint32_t importTableSize = directories[1]->get_size();

    if (importTableRVA == 0 || importTableSize == 0) {
      spdlog::error("Invalid import directory");
      return 0;
    }

    spdlog::info("Import directory at RVA: 0x{:x}, size: 0x{:x}", importTableRVA, importTableSize);

    // Find which section contains the import table
    COFFI::section *importSection = nullptr;
    uint32_t importSectionOffset = 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 (importTableRVA >= sectionRVA && importTableRVA < sectionRVA + sectionSize) {
        importSection = section;
        importSectionOffset = importTableRVA - sectionRVA;
        break;
      }
    }

    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 + importSectionOffset;

    // 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 lookupTableRVA = desc[i].originalFirstThunk;
      uint32_t iatRVA = desc[i].firstThunk;

      if (lookupTableRVA == 0)
        lookupTableRVA = iatRVA; // Use IAT if no lookup table

      const uint32_t *lookupTable =
          reinterpret_cast<const uint32_t *>(getRVAData(lookupTableRVA));
      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 iatEntryRVA = iatRVA + (j * sizeof(uint32_t));
          uint64_t iatEntryVA = imageBase + iatEntryRVA;
          spdlog::info("Found import '{}' in {} at IAT RVA: 0x{:x} (VA: 0x{:x})",
                       functionName, dllName, iatEntryRVA, iatEntryVA);
          return iatEntryVA;
        }
      }
    }

    spdlog::error("Could not find import: {}", functionName);
    return 0;
  }

  bool processRelocations(uint32_t injectedCodeRVA, uint32_t injectedCodeSize, 
                         COFFI::section *targetTextSection, COFFI::section *targetRdataSection,
                         COFFI::section *objMainSection, uint32_t &rdataUsedBytes) {
    auto &symbols = *objReader.get_symbols();
    auto &objSections = objReader.get_sections();
    auto &sectionRelocations = objMainSection->get_relocations();
    
    char *textSectionData = const_cast<char *>(targetTextSection->get_data());
    char *rdataSectionData = const_cast<char *>(targetRdataSection->get_data());
    uint32_t rdataAvailableSpace = findAvailableSpaceAtEnd(targetRdataSection);
    spdlog::info("Found {} bytes of available space in .rdata section", rdataAvailableSpace);
    
    uint32_t rdataStartOffset = targetRdataSection->get_data_size() - rdataAvailableSpace;
    char *rdataPtr = rdataSectionData + rdataStartOffset;
    const char *rdataEndPtr = rdataSectionData + targetRdataSection->get_data_size();

    spdlog::info("Processing relocations for injected code at RVA 0x{:x}...", injectedCodeRVA);

    for (auto &reloc : sectionRelocations) {
      if (reloc.get_type() != 6)
        continue; // Only handle type 6 (IMAGE_REL_I386_DIR32)

      uint32_t relocOffsetInCode = reloc.get_virtual_address();
      uint32_t relocRVA = injectedCodeRVA + relocOffsetInCode;
      std::string symbolName = reloc.get_symbol();
      uint64_t resolvedVA = 0;

      spdlog::info("Processing relocation at code offset 0x{:x} (RVA 0x{:x}, VA 0x{:x}) for symbol: {}",
                   relocOffsetInCode, relocRVA, imageBase + relocRVA, symbolName);

      if (symbolName.starts_with("__imp__")) {
        // Import symbol handling
        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);
        resolvedVA = findImportVA(functionName);

      } 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 = objSections[constSymbol->get_section_number() - 1];
          auto constData = constSection->get_data();
          uint32_t constOffset = constSymbol->get_value();

          // Find string length
          uint32_t stringLen = strlen(constData + constOffset);
          stringLen++; // Include null terminator

          if (rdataPtr + stringLen > rdataEndPtr) {
            spdlog::error("Not enough space in .rdata for string constant");
            return false;
          }

          // Copy string to .rdata
          std::memcpy(rdataPtr, constData + constOffset, stringLen);

          uint32_t stringRVA = (rdataPtr - rdataSectionData) + targetRdataSection->get_virtual_address();
          resolvedVA = imageBase + stringRVA;

          spdlog::info("Copied string constant '{}' to .rdata at RVA 0x{:x} (VA: 0x{:x})",
                       std::string(constData + constOffset), stringRVA, resolvedVA);

          rdataPtr += stringLen;
          rdataUsedBytes += stringLen;
        }
      }

      // Apply the relocation
      if (resolvedVA != 0) {
        uint32_t patchOffset = relocRVA - targetTextSection->get_virtual_address();
        uint32_t *patchPtr = reinterpret_cast<uint32_t *>(textSectionData + patchOffset);
        *patchPtr = static_cast<uint32_t>(resolvedVA); // Store the resolved VA
        spdlog::info("Applied relocation: patched RVA 0x{:x} (VA 0x{:x}) with VA 0x{:x}",
                     relocRVA, imageBase + relocRVA, resolvedVA);
      }
    }

    return true;
  }

  bool redirectEntryPoint(uint32_t targetCodeRVA) {
    // Patch the entry point to jump to our injected code
    auto optionalHeader = peReader.get_optional_header();
    if (!optionalHeader) {
      spdlog::error("Could not get optional header for entry point patching");
      return false;
    }
    
    uint32_t entryPointRVA = optionalHeader->get_entry_point_address();
    uint64_t entryPointVA = imageBase + entryPointRVA;
    uint64_t targetCodeVA = imageBase + targetCodeRVA;
    
    spdlog::info("Entry point at RVA: 0x{:x} (VA: 0x{:x})", entryPointRVA, entryPointVA);
    spdlog::info("Target injection at RVA: 0x{:x} (VA: 0x{:x})", targetCodeRVA, targetCodeVA);
    
    // Find which section contains the entry point
    COFFI::section* entrySection = nullptr;
    uint32_t entrySectionOffset = 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 (entryPointRVA >= sectionRVA && entryPointRVA < sectionRVA + sectionSize) {
        entrySection = section;
        entrySectionOffset = entryPointRVA - sectionRVA;
        break;
      }
    }
    
    if (!entrySection) {
      spdlog::error("Could not find section containing entry point");
      return false;
    }
    
    // Calculate relative jump offset
    // jmp rel32 instruction: E9 xx xx xx xx (5 bytes)
    // offset = target_address - (current_address + 5)
    int32_t jumpOffset = static_cast<int32_t>(targetCodeVA - (entryPointVA + 5));
    
    spdlog::info("Patching entry point with jmp instruction (offset: 0x{:x})", 
                 static_cast<uint32_t>(jumpOffset));
    
    // Patch the entry point with jmp instruction
    char* entryData = const_cast<char*>(entrySection->get_data()) + entrySectionOffset;
    entryData[0] = static_cast<char>(0xE9); // jmp rel32 opcode
    *reinterpret_cast<int32_t*>(&entryData[1]) = jumpOffset;
    
    spdlog::info("Entry point patched successfully");
    return true;
  }

  bool patchAndSave() {
    // Check available space in both sections
    uint32_t textAvailableSpace = findAvailableSpaceAtEnd(textSection);

    spdlog::info("Found {} bytes of available space (null bytes) at end of .text section", textAvailableSpace);

    if (textAvailableSpace < mainSize) {
      spdlog::error("Not enough space in .text section! Need {} bytes, found {} bytes", mainSize, textAvailableSpace);
      return false;
    }

    // Calculate injection points
    uint32_t textSectionSize = textSection->get_data_size();
    uint32_t textInjectionOffset = textSectionSize - textAvailableSpace;
    uint32_t injectionRVA = textSection->get_virtual_address() + textInjectionOffset;

    spdlog::info("Injecting {} bytes at .text section offset 0x{:x} (RVA: 0x{:x}, VA: 0x{:x})",
                 mainSize, textInjectionOffset, injectionRVA, imageBase + injectionRVA);

    // Copy main function code to injection point
    char *textSectionData = const_cast<char *>(textSection->get_data());
    char *injectionPtr = textSectionData + textInjectionOffset;
    const uint8_t *mainCode = reinterpret_cast<const uint8_t *>(mainSection->get_data()) + mainOffset;
    std::memcpy(injectionPtr, mainCode, mainSize);

    spdlog::info("Injected code into existing .text section space (size unchanged: 0x{:x})", textSectionSize);

    // Process relocations
    uint32_t rdataUsedBytes = 0;
    if (!processRelocations(injectionRVA, mainSize, textSection, rdataSection, mainSection, rdataUsedBytes)) {
      return false;
    }

    spdlog::info("Used {} bytes in .rdata section for string constants", rdataUsedBytes);

    // Patch the entry point to jump to our injected code
    if (!redirectEntryPoint(injectionRVA)) {
      return false;
    }

    // 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 RVA: 0x{:x} (VA: 0x{:x})",
                 injectionRVA, imageBase + injectionRVA);
    return true;
  }

  bool run() {
    if (!loadFiles())
      return false;
    if (!findMainSymbol())
      return false;

    calculateMainSize();
    logMainFunctionCode(mainSection, mainOffset, mainSize);
    logRelocations(mainSection, mainSymbol->get_section_number());

    if (!validateSections())
      return false;
    if (!patchAndSave())
      return false;

    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[]) {
  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;
}