// Script to find hardcoded addresses in the binary that need to be relocated
// @category _Reman3
// @menupath Reman3.Find and dump Relocations

import ghidra.app.script.GhidraScript;
import ghidra.program.model.listing.*;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.program.model.address.*;
import ghidra.program.model.scalar.Scalar;
import ghidra.program.model.symbol.Reference;
import java.util.*;
import java.io.*;
import re3lib.RemanConfig;

public class FindRelocations extends GhidraScript {
  private Set<String> foundRelocations = new HashSet<>();
  private PrintWriter outputFile;

  long addrMin, addrMax;

  @Override
  public void run() throws Exception {
    RemanConfig.INSTANCE = new RemanConfig(this);
    RemanConfig.INSTANCE.createDirectories();

    addrMin = RemanConfig.INSTANCE.staticMemoryBlockStart.getOffset();
    addrMax = RemanConfig.INSTANCE.staticMemoryBlockEnd.getOffset();

    // Create output file for relocations
    File relocFile = new File(RemanConfig.INSTANCE.outputDir, "relocations.def");
    outputFile = new PrintWriter(new FileWriter(relocFile));

    try {
      println("Counting instructions and data for progress tracking...");

      // Count total work for progress monitoring
      long totalInstructions = currentProgram.getListing().getNumInstructions();
      long totalDataBytes = currentProgram.getMemory().getLoadedAndInitializedAddressSet().getNumAddresses();
      long totalWork = totalInstructions + (totalDataBytes / currentProgram.getDefaultPointerSize());

      monitor.initialize(totalWork);
      monitor.setMessage("Scanning for relocations...");

      println("Scanning " + totalInstructions + " instructions and " + totalDataBytes + " data bytes...");

      // Get all instructions in the program
      InstructionIterator instructions = currentProgram.getListing().getInstructions(true);
      long processedInstructions = 0;

      while (instructions.hasNext()) {
        if (monitor.isCancelled()) {
          println("Operation cancelled by user");
          return;
        }

        Instruction instruction = instructions.next();
        analyzeInstruction(instruction);

        // Check if we've gone past the end address
        if (instruction.getAddress().getOffset() >= 0x00844190) {
          break;
        }

        processedInstructions++;
        if (processedInstructions % 1000 == 0) {
          monitor.setProgress(processedInstructions);
          monitor.setMessage("Processed " + processedInstructions + "/" + totalInstructions + " instructions...");
        }
      }

      monitor.setProgress(totalInstructions);
      monitor.setMessage("Analyzing data sections...");

      // Also check data references
      analyzeDataReferences(totalInstructions, totalWork);

      monitor.setMessage("Scan complete");
      println("Found " + foundRelocations.size() + " potential relocations");
      println("Results saved to: " + relocFile.getAbsolutePath());

    } finally {
      if (outputFile != null) {
        outputFile.close();
      }
    }
  }

  private void analyzeInstruction(Instruction instruction) {
    String mnemonic = instruction.getMnemonicString().toLowerCase();

    // Check for instructions that commonly use absolute addresses
    if (isRelocatableInstruction(mnemonic)) {
      // Check references from this instruction
      Reference[] refs = instruction.getReferencesFrom();
      for (Reference ref : refs) {
        Address toAddr = ref.getToAddress();
        
        if (isInMainMemorySpace(toAddr)) {
          // Check if the target address appears in the instruction bytes (absolute addressing)
          int operandOffset = findAbsoluteAddressOffset(instruction, toAddr);
          if (operandOffset >= 0) {
            recordRelocation(instruction.getAddress(), toAddr, mnemonic, 
                           "absolute_" + ref.getReferenceType().getName(), operandOffset);
          }
        }
      }
    }
  }

  private int findAbsoluteAddressOffset(Instruction instruction, Address targetAddr) {
    try {
      byte[] instructionBytes = instruction.getBytes();
      long targetValue = targetAddr.getOffset();
      
      // Convert target address to little-endian byte array (x86 32-bit)
      byte[] targetBytes = new byte[4];
      targetBytes[0] = (byte) (targetValue & 0xFF);
      targetBytes[1] = (byte) ((targetValue >> 8) & 0xFF);
      targetBytes[2] = (byte) ((targetValue >> 16) & 0xFF);
      targetBytes[3] = (byte) ((targetValue >> 24) & 0xFF);
      
      // Search for the target address bytes in the instruction and return offset
      int offset = findSequenceOffset(instructionBytes, targetBytes);
      
      return offset;
      
    } catch (Exception e) {
      return -1;
    }
  }
  
  private String bytesToHex(byte[] bytes) {
    StringBuilder sb = new StringBuilder();
    for (int i = 0; i < bytes.length; i++) {
      if (i > 0) sb.append(" ");
      sb.append(String.format("%02x", bytes[i] & 0xFF));
    }
    return sb.toString();
  }

  private int findSequenceOffset(byte[] haystack, byte[] needle) {
    if (needle.length > haystack.length) {
      return -1;
    }
    
    for (int i = 0; i <= haystack.length - needle.length; i++) {
      boolean found = true;
      for (int j = 0; j < needle.length; j++) {
        if (haystack[i + j] != needle[j]) {
          found = false;
          break;
        }
      }
      if (found) {
        return i; // Return the offset where the sequence starts
      }
    }
    return -1;
  }

  private boolean isRelocatableInstruction(String mnemonic) {
    // Instructions that commonly use absolute addresses
    return mnemonic.equals("mov") || mnemonic.equals("lea") ||
           mnemonic.equals("call") || mnemonic.equals("jmp") ||
           mnemonic.equals("push") || mnemonic.equals("cmp") ||
           mnemonic.equals("test") || mnemonic.equals("add") ||
           mnemonic.equals("sub") || mnemonic.equals("and") ||
           mnemonic.equals("or") || mnemonic.equals("xor") ||
           mnemonic.startsWith("j"); // All jumps, we'll filter by byte analysis
  }

  private void analyzeDataReferences(long instructionsProcessed, long totalWork) {
    // Only scan actual data sections, not code sections
    MemoryBlock[] blocks = currentProgram.getMemory().getBlocks();
    long processedBytes = 0;
    int pointerSize = currentProgram.getDefaultPointerSize();
    
    for (MemoryBlock block : blocks) {
      // Skip executable blocks (code sections)
      if (block.isExecute()) {
        continue;
      }
      
      // Only scan initialized data blocks
      if (!block.isInitialized()) {
        continue;
      }
      
      if (monitor.isCancelled()) {
        println("Operation cancelled by user");
        return;
      }

      if (block.getName().equals(".rsrc")) {
        continue;
      }
      
      println("Scanning data block: " + block.getName() + " (" + block.getStart() + " - " + block.getEnd() + ")");
      
      Address addr = block.getStart();
      while (addr != null && addr.compareTo(block.getEnd()) <= 0) {
        try {
          // Check if this location contains a pointer-sized value
          byte[] bytes = new byte[pointerSize];
          currentProgram.getMemory().getBytes(addr, bytes);

          long value = 0;
          for (int i = 0; i < bytes.length; i++) {
            value |= ((long) (bytes[i] & 0xFF)) << (i * 8);
          }

          if (looksLikeAddress(value)) {
            try {
              Address targetAddr = currentProgram.getAddressFactory().getDefaultAddressSpace().getAddress(value);
              if (isInMainMemorySpace(targetAddr)) {
                recordRelocation(addr, targetAddr, "data", "pointer", 0);
              }
            } catch (Exception e) {
              // Invalid address, ignore
            }
          }

          addr = addr.add(pointerSize);
          processedBytes += pointerSize;

        } catch (Exception e) {
          addr = addr.add(1);
          processedBytes++;
        }
      }
    }
  }

  private boolean isInMainMemorySpace(Address addr) {
    if (addr == null)
      return false;

    // Check if address is in a loaded memory block
    return currentProgram.getMemory().contains(addr) &&
        !addr.getAddressSpace().isOverlaySpace();
  }

  private boolean looksLikeAddress(long value) {
    // Heuristic: check if value is in reasonable address range
    // Adjust these ranges based on your target architecture
    return value >= addrMin
        && value <= addrMax; // Typical executable range
  }

  private void recordRelocation(Address fromAddr, Address toAddr, String instruction, String type, int operandOffset) {
    String relocKey = fromAddr.toString() + " -> " + toAddr.toString();
    
    if (foundRelocations.add(relocKey)) {
      String instructionBytes = getInstructionBytesString(fromAddr);
      Address operandPtr = fromAddr.add(operandOffset);
      
      String line = String.format("REL(0x%s, 0x%s, 0x%s) // %s [%s] | %s",
          fromAddr.toString(),
          operandPtr.toString(),
          toAddr.toString(),
          instruction,
          type,
          instructionBytes);
      
      outputFile.println(line);
      outputFile.flush();
    }
  }

  private String getInstructionBytesString(Address addr) {
    try {
      Instruction instruction = currentProgram.getListing().getInstructionAt(addr);
      if (instruction != null) {
        byte[] bytes = instruction.getBytes();
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < bytes.length; i++) {
          if (i > 0)
            sb.append(" ");
          sb.append(String.format("%02x", bytes[i] & 0xFF));
        }
        return sb.toString();
      } else {
        // This is a data reference, show the pointer bytes
        byte[] bytes = new byte[4]; // Show 4 bytes for data
        int bytesRead = currentProgram.getMemory().getBytes(addr, bytes);
        if (bytesRead > 0) {
          StringBuilder sb = new StringBuilder();
          for (int i = 0; i < bytesRead; i++) {
            if (i > 0) sb.append(" ");
            sb.append(String.format("%02x", bytes[i] & 0xFF));
          }
          return sb.toString();
        }
      }
    } catch (Exception e) {
      // Check if this is in an initialized memory block
      if (currentProgram.getMemory().contains(addr)) {
        return "unreadable";
      }
    }
    return "??";
  }

  private boolean isPartOfFunction(Address addr) {
    // Check if there's an instruction at this address
    Instruction instruction = currentProgram.getListing().getInstructionAt(addr);
    if (instruction != null) {
      return true;
    }
    
    // Check if this address is within any function's body
    Function function = currentProgram.getFunctionManager().getFunctionContaining(addr);
    if (function != null) {
      // Additional check: make sure we're in the function body, not just data referenced by it
      AddressSetView functionBody = function.getBody();
      return functionBody.contains(addr);
    }
    
    return false;
  }
}