// 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;
import ghidra.util.task.TaskMonitor;
import ghidra.program.model.pcode.PcodeOp;
import ghidra.program.model.pcode.Varnode;

public class FindRelocations extends GhidraScript {
  private Set<Address> foundRelocations = new HashSet<>();
  private PrintWriter outputFile;
  
  @Override
  public void run() throws Exception {
    RemanConfig.INSTANCE = new RemanConfig(this);
    RemanConfig.INSTANCE.createDirectories();
    
    // 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();
    
    // Use pcode analysis for better accuracy
    if (isRelocatableInstruction(mnemonic)) {
      analyzeInstructionPcode(instruction);
      
      // Still check operands for direct absolute addresses
      for (int i = 0; i < instruction.getNumOperands(); i++) {
        analyzeOperand(instruction, i);
      }
    }
  }
  
  private void analyzeInstructionPcode(Instruction instruction) {
    try {
      PcodeOp[] pcode = instruction.getPcode();
      
      for (PcodeOp op : pcode) {
        // Look for operations that use absolute addresses
        switch (op.getOpcode()) {
          case PcodeOp.LOAD:
          case PcodeOp.STORE:
            // Memory operations - check if using absolute address
            analyzeMemoryOperation(instruction, op);
            break;
            
          case PcodeOp.CALL:
          case PcodeOp.CALLIND:
            // Call operations - check if target is absolute
            analyzeCallOperation(instruction, op);
            break;
            
          case PcodeOp.BRANCH:
          case PcodeOp.BRANCHIND:
            // Branch operations - usually relative, but check indirect branches
            analyzeBranchOperation(instruction, op);
            break;
            
          case PcodeOp.COPY:
            // Copy operations that move absolute addresses
            analyzeCopyOperation(instruction, op);
            break;
        }
      }
    } catch (Exception e) {
      // Fallback to reference analysis if pcode fails
      analyzeReferences(instruction);
    }
  }
  
  private void analyzeMemoryOperation(Instruction instruction, PcodeOp op) {
    // LOAD/STORE operations with absolute addresses
    Varnode input = op.getInput(1); // Address input
    if (input != null && input.isConstant()) {
      long addr = input.getOffset();
      if (looksLikeAddress(addr)) {
        try {
          Address targetAddr = currentProgram.getAddressFactory().getDefaultAddressSpace().getAddress(addr);
          if (isInMainMemorySpace(targetAddr)) {
            recordRelocation(instruction.getAddress(), targetAddr, 
                           instruction.getMnemonicString(), "memory_op");
          }
        } catch (Exception e) {
          // Invalid address
        }
      }
    }
  }
  
  private void analyzeCallOperation(Instruction instruction, PcodeOp op) {
    // Only flag indirect calls or calls with absolute targets
    if (op.getOpcode() == PcodeOp.CALLIND) {
      // Indirect call - check if target address is absolute
      Varnode target = op.getInput(0);
      if (target != null && target.isConstant()) {
        long addr = target.getOffset();
        if (looksLikeAddress(addr)) {
          try {
            Address targetAddr = currentProgram.getAddressFactory().getDefaultAddressSpace().getAddress(addr);
            if (isInMainMemorySpace(targetAddr)) {
              recordRelocation(instruction.getAddress(), targetAddr, 
                             instruction.getMnemonicString(), "indirect_call");
            }
          } catch (Exception e) {
            // Invalid address
          }
        }
      }
    }
    // Direct calls are usually relative, skip them
  }
  
  private void analyzeBranchOperation(Instruction instruction, PcodeOp op) {
    // Only flag indirect branches (direct branches are relative)
    if (op.getOpcode() == PcodeOp.BRANCHIND) {
      Varnode target = op.getInput(0);
      if (target != null && target.isConstant()) {
        long addr = target.getOffset();
        if (looksLikeAddress(addr)) {
          try {
            Address targetAddr = currentProgram.getAddressFactory().getDefaultAddressSpace().getAddress(addr);
            if (isInMainMemorySpace(targetAddr)) {
              recordRelocation(instruction.getAddress(), targetAddr, 
                             instruction.getMnemonicString(), "indirect_branch");
            }
          } catch (Exception e) {
            // Invalid address
          }
        }
      }
    }
    // Direct branches (BRANCH) are relative, ignore them
  }
  
  private void analyzeCopyOperation(Instruction instruction, PcodeOp op) {
    // COPY operations that load absolute addresses (like LEA or MOV immediate)
    Varnode input = op.getInput(0);
    if (input != null && input.isConstant()) {
      long addr = input.getOffset();
      if (looksLikeAddress(addr)) {
        try {
          Address targetAddr = currentProgram.getAddressFactory().getDefaultAddressSpace().getAddress(addr);
          if (isInMainMemorySpace(targetAddr)) {
            recordRelocation(instruction.getAddress(), targetAddr, 
                           instruction.getMnemonicString(), "load_address");
          }
        } catch (Exception e) {
          // Invalid address
        }
      }
    }
  }
  
  private void analyzeReferences(Instruction instruction) {
    // Fallback method - but be more selective
    Reference[] refs = instruction.getReferencesFrom();
    for (Reference ref : refs) {
      if (ref.getReferenceType().isData()) {
        // Data references are usually absolute
        Address toAddr = ref.getToAddress();
        if (isInMainMemorySpace(toAddr)) {
          recordRelocation(instruction.getAddress(), toAddr, 
                         instruction.getMnemonicString(), "data_reference");
        }
      }
      // Skip flow references (calls/jumps) as they're usually relative
    }
  }
  
  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 >= 0x400000 && value <= 0x7FFFFFFF; // Typical executable range
  }
  
  private void recordRelocation(Address fromAddr, Address toAddr, String instruction, String type) {
    if (foundRelocations.add(toAddr)) {
      String line = String.format("0x%s -> 0x%s (%s) [%s]", 
                                 fromAddr.toString(), 
                                 toAddr.toString(), 
                                 instruction, 
                                 type);
      println(line);
      outputFile.println(line);
      outputFile.flush();
    }
  }
}