Code Patcher

patcher/patcher.cpp

@@ -7,8 +7,10 @@ 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();
+  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);
 
@@ -49,12 +51,49 @@ int main(int argc, char *argv[]) {
   auto &sections = objReader.get_sections();
   auto mainSection = sections[mainSymbol->get_section_number() - 1];
 
-  // Calculate main function size and get its code
-  uint32_t mainSize = mainSymbol->get_value() + mainSymbol->get_auxiliary_symbols().size(); // This needs proper calculation
-  auto mainCodeData = mainSection->get_data();
+  // Calculate main function size using next symbol method
   uint32_t mainOffset = mainSymbol->get_value();
+  uint32_t mainSize = 0;
 
-  SPDLOG_INFO("Found main function at offset {} with estimated size {}", mainOffset, mainSize);
+  // 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();
@@ -71,8 +110,10 @@ int main(int argc, char *argv[]) {
     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);
+  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;
 }