From b1dc3a13c2c068f8f560d181a8d1f101647196ed Mon Sep 17 00:00:00 2001
From: Guus Waals <_@guusw.nl>
Date: Fri, 6 Jun 2025 17:49:59 +0800
Subject: [PATCH] Code Patcher

---
 patcher/patcher.cpp | 73 +++++++++++++++++++++++++++++++++++----------
 1 file changed, 57 insertions(+), 16 deletions(-)

diff --git a/patcher/patcher.cpp b/patcher/patcher.cpp
index 79ddcf4e..98b9b911 100644
--- a/patcher/patcher.cpp
+++ b/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);
 
@@ -30,8 +32,8 @@ int main(int argc, char *argv[]) {
   }
 
   // Find the 'main' function in the object file
-  auto& symbols = *objReader.get_symbols();
-  COFFI::symbol* mainSymbol = nullptr;
+  auto &symbols = *objReader.get_symbols();
+  COFFI::symbol *mainSymbol = nullptr;
   for (auto &sym : symbols) {
     SPDLOG_INFO("Symbol: {}", sym.get_name());
     if (sym.get_name() == "_ref") {
@@ -46,20 +48,57 @@ int main(int argc, char *argv[]) {
   }
 
   // Get the section containing the main function
-  auto& sections = objReader.get_sections();
+  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();
-  
-  SPDLOG_INFO("Found main function at offset {} with estimated size {}", mainOffset, mainSize);
+  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) {
+  auto &peSections = peReader.get_sections();
+  COFFI::section *textSection = nullptr;
+  for (auto &section : peSections) {
     if (section->get_name() == ".text") {
       textSection = section;
       break;
@@ -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;
 }
\ No newline at end of file