// 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.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<Address> 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.txt");
    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);

        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 operands for absolute addresses
      for (int i = 0; i < instruction.getNumOperands(); i++) {
        analyzeOperand(instruction, i);
      }

      // Check references from this instruction - but filter out relative references
      Reference[] refs = instruction.getReferencesFrom();
      for (Reference ref : refs) {
        // Skip relative references (jumps/calls with relative addressing)
        if (ref.getReferenceType().isCall() || ref.getReferenceType().isJump()) {
          // For jumps and calls, check if it's using absolute addressing
          if (usesAbsoluteAddressing(instruction, ref)) {
            Address toAddr = ref.getToAddress();
            if (isInMainMemorySpace(toAddr)) {
              recordRelocation(instruction.getAddress(), toAddr, mnemonic, "absolute_reference");
            }
          }
        } else if (ref.getReferenceType().isData()) {
          // Data references are more likely to be absolute
          Address toAddr = ref.getToAddress();
          if (isInMainMemorySpace(toAddr)) {
            recordRelocation(instruction.getAddress(), toAddr, mnemonic, "data_reference");
          }
        }
      }
    }
  }

  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"); // conditional jumps
  }

  private void analyzeOperand(Instruction instruction, int opIndex) {
    Object[] operandObjects = instruction.getOpObjects(opIndex);

    for (Object obj : operandObjects) {
      if (obj instanceof Address) {
        Address addr = (Address) obj;
        if (isInMainMemorySpace(addr)) {
          recordRelocation(instruction.getAddress(), addr,
              instruction.getMnemonicString(), "operand_" + opIndex);
        }
      } else if (obj instanceof Scalar) {
        Scalar scalar = (Scalar) obj;
        // Check if scalar value looks like an address in our memory space
        long value = scalar.getUnsignedValue();
        if (looksLikeAddress(value)) {
          try {
            Address addr = currentProgram.getAddressFactory().getDefaultAddressSpace().getAddress(value);
            if (isInMainMemorySpace(addr)) {
              recordRelocation(instruction.getAddress(), addr,
                  instruction.getMnemonicString(), "scalar_" + opIndex);
            }
          } catch (Exception e) {
            // Invalid address, ignore
          }
        }
      }
    }
  }

  private void analyzeDataReferences(long instructionsProcessed, long totalWork) {
    // Check data sections for absolute addresses
    AddressSetView dataAddresses = currentProgram.getMemory().getLoadedAndInitializedAddressSet();
    long totalDataBytes = dataAddresses.getNumAddresses();
    long processedBytes = 0;
    int pointerSize = currentProgram.getDefaultPointerSize();

    for (AddressRange range : dataAddresses) {
      if (monitor.isCancelled()) {
        println("Operation cancelled by user");
        return;
      }

      Address addr = range.getMinAddress();
      while (addr != null && addr.compareTo(range.getMaxAddress()) <= 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");
              }
            } catch (Exception e) {
              // Invalid address, ignore
            }
          }

          addr = addr.add(pointerSize); // Jump by pointer size for efficiency
          processedBytes += pointerSize;

          // Update progress every 10000 bytes
          if (processedBytes % 10000 == 0) {
            long currentProgress = instructionsProcessed + (processedBytes / pointerSize);
            monitor.setProgress(currentProgress);
            monitor.setMessage("Analyzing data: " + (processedBytes * 100 / totalDataBytes) + "% complete");
          }

        } 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 boolean usesAbsoluteAddressing(Instruction instruction, Reference ref) {
    // Check the instruction bytes to determine addressing mode
    String mnemonic = instruction.getMnemonicString().toLowerCase();

    // For x86, most conditional jumps (JA, JE, JNE, etc.) use relative addressing
    if (mnemonic.startsWith("j") && !mnemonic.equals("jmp")) {
      return false; // Conditional jumps are typically relative
    }

    // For JMP and CALL, check the operand representation
    for (int i = 0; i < instruction.getNumOperands(); i++) {
      String operandStr = instruction.getDefaultOperandRepresentation(i);
      // If operand shows as a direct address (not offset), it might be absolute
      // But we need to be more sophisticated here...

      // Check if this is an indirect reference [address] which would be absolute
      if (operandStr.contains("[") && operandStr.contains("]")) {
        return true; // Indirect addressing typically uses absolute addresses
      }
    }

    // For now, assume most jumps/calls are relative unless proven otherwise
    return false;
  }

  private void recordRelocation(Address fromAddr, Address toAddr, String instruction, String type) {
    if (foundRelocations.add(toAddr)) {
      String instructionBytes = getInstructionBytesString(fromAddr);
      String line = String.format("0x%s -> 0x%s (%s) [%s] | %s",
          fromAddr.toString(),
          toAddr.toString(),
          instruction,
          type,
          instructionBytes);
      println(line);
      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();
      }
    } catch (Exception e) {
      // If it's a data reference, try to get the bytes at that location
      try {
        byte[] bytes = new byte[4]; // Show 4 bytes for data
        currentProgram.getMemory().getBytes(addr, 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();
      } catch (Exception ex) {
        return "??";
      }
    }
    return "??";
  }
}