WIP Patch

patcher/patcher.cpp

@@ -117,7 +117,8 @@ struct Patcher {
       ptr--;
     }
 
-    return availableSpace;
+    // Keep 1 byte marging, might be another string or terminator
+    return availableSpace - 1;
   }
 
   void logMainFunctionCode() {
@@ -181,16 +182,111 @@ struct Patcher {
     return true;
   }
 
-  // Tries to find an import in the input PE file
+  // Tries to find an import in the input PE file and return its IAT address
-  void findImportInPE(std::string_view name) {
+  uint32_t findImportInPE(const std::string &functionName) {
-    auto &imports = peReader.get_imports();
+    auto &directories = peReader.get_directories();
-    for (auto &import : imports) {
+
-      if (import.get_name() == name) {
+    // Import table is directory entry 1
-        spdlog::info("Found import: {}", name);
+    if (directories.size() <= 1 || !directories[1]) {
-        return;
+      spdlog::error("No import directory found in PE file");
+      return 0;
+    }
+
+    uint32_t importRVA = directories[1]->get_virtual_address();
+    uint32_t importSize = directories[1]->get_size();
+
+    if (importRVA == 0 || importSize == 0) {
+      spdlog::error("Invalid import directory");
+      return 0;
+    }
+
+    spdlog::info("Import directory at RVA: 0x{:x}, size: 0x{:x}", importRVA,
+                 importSize);
+
+    // Find which section contains the import table
+    COFFI::section *importSection = nullptr;
+    uint32_t importFileOffset = 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 (importRVA >= sectionRVA && importRVA < sectionRVA + sectionSize) {
+        importSection = section;
+        importFileOffset = importRVA - sectionRVA;
+        break;
       }
     }
-    spdlog::error("Could not find import: {}", name);
+
+    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 + importFileOffset;
+
+    // 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 lookupRVA = desc[i].originalFirstThunk;
+      uint32_t iatRVA = desc[i].firstThunk;
+
+      if (lookupRVA == 0)
+        lookupRVA = iatRVA; // Use IAT if no lookup table
+
+      const uint32_t *lookupTable =
+          reinterpret_cast<const uint32_t *>(getRVAData(lookupRVA));
+      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 iatAddress = imageBase + iatRVA + (j * sizeof(uint32_t));
+          spdlog::info("Found import '{}' in {} at IAT address: 0x{:x}",
+                       functionName, dllName, iatAddress);
+          return iatAddress;
+        }
+      }
+    }
+
+    spdlog::error("Could not find import: {}", functionName);
+    return 0;
   }
 
   bool processRelocations(char *textStart, char *textEnd, char *&rdataPtr) {
@@ -223,18 +319,15 @@ struct Patcher {
       spdlog::info("Processing relocation at VA 0x{:x} for symbol: {}",
                    relocRVA + imageBase, symbolName);
 
-      if (symbolName.starts_with("__imp_")) {
+      if (symbolName.starts_with("__imp__")) {
         // Import symbol handling
-        std::string functionName = symbolName.substr(6);
+        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);
-        spdlog::warn("Import resolution not fully implemented yet for: {}",
+        resolvedAddress = findImportInPE(functionName);
-                     functionName);
-        resolvedAddress = 0x12345678; // Placeholder
 
       } else if (symbolName.starts_with("??_C@")) {
         // String constant handling
@@ -252,10 +345,7 @@ struct Patcher {
           uint32_t constOffset = constSymbol->get_value();
 
           // Find string length
-          uint32_t stringLen = 0;
+          uint32_t stringLen = strlen(constData + constOffset);
-          while (constData[constOffset + stringLen] != 0) {
-            stringLen++;
-          }
           stringLen++; // Include null terminator
 
           if (rdataPtr + stringLen > rdataEndPtr) {
@@ -287,7 +377,7 @@ struct Patcher {
         ptr[0] = resolvedAddress;
         spdlog::info(
             "Applied relocation: patched offset 0x{:x} with address 0x{:x}",
-            patchOffset, resolvedAddress);
+            relocRVA + imageBase, resolvedAddress);
       }
     }
 
@@ -339,16 +429,16 @@ struct Patcher {
         reinterpret_cast<const uint8_t *>(mainSection->get_data()) + mainOffset;
     std::memcpy(dstText, mainCode, mainSize);
 
+    spdlog::info("Injected code into existing .text section space (size "
+                 "unchanged: 0x{:x})",
+                 textSize);
+
     // Process relocations
     char *rdataPtr = dstRdata;
     if (!processRelocations(dstText, dstText + mainSize, rdataPtr)) {
       return false;
     }
 
-    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)) {
@@ -379,6 +469,28 @@ struct Patcher {
 
     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[]) {