// Script to sanitize global symbols in Ghidra
// @category _Reman3

import java.io.File;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.List;

import ghidra.app.cmd.label.AddLabelCmd;
import ghidra.app.decompiler.DecompileResults;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.VariableStorage;
import ghidra.program.model.pcode.HighSymbol;
import ghidra.program.model.symbol.SourceType;
import ghidra.program.model.symbol.Symbol;
import re3lib.RecompileConfig;
import re3lib.Utils;

public class SanitizeGlobalSymbols extends GhidraScript {
  private static final boolean AUTO_RENAME_SYMBOLS = true;

  HashSet<Address> functionAddrBlackList = new HashSet<>();

  boolean shouldDecompileFunction(Function function) {
    return !functionAddrBlackList.contains(function.getEntryPoint());
  }

  void sanitizeGlobalSymbolsPass(List<Function> functions) {
    Hashtable<String, HighSymbol> globalSymbols = new Hashtable<>();

    for (Function function : functions) {
      println("Processing global symbols for " + function.getName());

      DecompileResults decompRes = RecompileConfig.INSTANCE.decompCache.getOrInsert(function);
      Iterator<HighSymbol> smyIt = decompRes.getHighFunction().getGlobalSymbolMap().getSymbols();

      HighSymbol gsym = smyIt.next();
      if (globalSymbols.containsKey(gsym.getName()))
        continue;
      println("GLOBAL: " + gsym.getName());
      String sanitizedName = sanitizeFunctionName(gsym.getName());
      if (!sanitizedName.equals(gsym.getName())) {
        if (AUTO_RENAME_SYMBOLS) {
          Symbol symbol = gsym.getSymbol();
          VariableStorage storage = gsym.getStorage();
          Address addr = storage.getMinAddress();
          println("Renaming global symbol: " + gsym.getName() + " (" + addr
              + ") -> " + sanitizedName);
          if (symbol != null) {
            AddLabelCmd cmd = new AddLabelCmd(addr, sanitizedName, symbol.getParentNamespace(),
                SourceType.USER_DEFINED);
            if (cmd.applyTo(currentProgram)) {
              println("Renamed global symbol: " + gsym.getName() + " -> " + sanitizedName);
            } else {
              println("Error renaming symbol: " + cmd.getStatusMsg());
            }
          } else {
            println("Symbol is null: " + gsym.getName() + " - " + function.getName());
          }
        } else {
          println("Invalid global symbol name: " + gsym.getName() + " - " + function.getName());
        }
      }
    }
  }

  private static String sanitizeFunctionName(String name) {
    return name.replaceAll("[^a-zA-Z0-9_]", "_");
  }

  @Override
  public void run() throws Exception {
    if (currentProgram == null) {
      return;
    }

    RecompileConfig.INSTANCE = new RecompileConfig(this);

    if (!new File(RecompileConfig.INSTANCE.outputDir).exists()) {
      throw new Exception("Output directory does not exist: " + RecompileConfig.INSTANCE.outputDir);
    }

    functionAddrBlackList = Utils.loadFunctionBlacklist(RecompileConfig.INSTANCE.functionBlacklistPath);

    List<Function> functions = new ArrayList<>();

    Iterator<Function> functionsIt = currentProgram.getFunctionManager().getFunctions(true).iterator();
    while (functionsIt.hasNext()) {
      Function function = functionsIt.next();
      if (!shouldDecompileFunction(function)) {
        continue;
      }

      functions.add(function);
    }

    sanitizeGlobalSymbolsPass(functions);
  }
}