reman3/GhidraScripts/re3lib/TypeDumper.java

package re3lib;

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

import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.data.CategoryPath;
import ghidra.program.model.data.Composite;
import ghidra.program.model.data.DataType;
import ghidra.program.model.data.Enum;
import ghidra.program.model.data.EnumDataType;
import ghidra.program.model.data.FunctionDefinition;
import ghidra.program.model.data.Pointer;
import ghidra.program.model.data.ProgramBasedDataTypeManager;
import ghidra.program.model.data.Structure;
import ghidra.program.model.data.TypeDef;
import ghidra.program.model.data.TypedefDataType;
import ghidra.program.model.data.Union;
import ghidra.program.model.data.UnionDataType;
import ghidra.program.model.listing.Data;
import ghidra.program.model.listing.Program;
import re3lib.FunctionDatabase.GlobalEntry;

public class TypeDumper {
  Program currentProgram;
  GhidraScript script;

  public TypeDumper(GhidraScript script) {
    this.script = script;
    currentProgram = script.getCurrentProgram();
    RemanConfig.INSTANCE = new RemanConfig(script);
  }

  public void run() throws Exception {
    ProgramBasedDataTypeManager dtm = currentProgram.getDataTypeManager();

    HashSet<String> requiredFunctionTypes = new HashSet<>();
    FunctionDatabase fd = new FunctionDatabase(script);
    List<GlobalEntry> globals = fd.loadAllGlobals();
    for (GlobalEntry globalEntry : globals) {
      // Fetch type of global from ghidra
      Address addr = globalEntry.address;
      if (globalEntry.name.equals("p_fn_vDisplayAll"))
        script.println("EE" + globalEntry.name + "EE");
      // script.print("Global: " + globalEntry.name);

      Data data = script.getDataAt(addr);
      if (data != null) {
        DataType type = data.getDataType();
        if (type instanceof Pointer) {
          DataType inner = ((Pointer) type).getDataType();
          if (inner instanceof FunctionDefinition) {
            requiredFunctionTypes.add(inner.getDisplayName());
            script.println("Adding required global function type: " + type.getDisplayName());
          }
        }
      }
    }

    HashSet<String> typeBlacklist = Utils.loadSimpleBlacklist(RemanConfig.INSTANCE.typeBlacklistPath);
    HashSet<String> categoryPathBlacklist = Utils
        .loadSimpleBlacklist(RemanConfig.INSTANCE.categoryPathBlacklistPath);

    if (typeBlacklist == null) {
      script.println("Building struct blacklist from existing data types");
      typeBlacklist = new HashSet<>();
      Iterator<DataType> it = dtm.getAllDataTypes();
      while (it.hasNext()) {
        DataType dt = it.next();
        if (dt instanceof Structure || dt instanceof TypedefDataType) {
          typeBlacklist.add(dt.getDisplayName());
        }
      }
      Utils.saveStructBlacklist(typeBlacklist, RemanConfig.INSTANCE.typeBlacklistPath);
    }

    List<DataType> filteredTypes = new ArrayList<>();

    // Iterator<DataType> compIt = dtm.getAllDataTypes();
    // while (compIt.hasNext()) {
    // DataType comp = compIt.next();
    // // script.println("Found: " + comp.getDisplayName() + " - " +
    // // comp.getClass().getSimpleName());

    // if (comp instanceof TypeDef) {
    // if (comp.getDisplayName().startsWith("FIL_")) {
    // script.println("Found: " + comp.getDisplayName() + " - " + comp.getName() + "
    // - " + comp.getClass().getSimpleName());
    // }
    // }

    // // if (comp.getName() == "FIL_tdstConcatFile") {
    // // // script.println("Found: " + dt.getDisplayName() + " - " +
    // // // dt.getClass().getSimpleName());
    // // throw new Exception("Found: " + comp.getDisplayName() + " - " +
    // comp.getClass().getSimpleName());
    // // }
    // }

    Iterator<DataType> it = dtm.getAllDataTypes();
    while (it.hasNext()) {
      DataType dt = it.next();

      if (typeBlacklist.contains(dt.getDisplayName()))
        continue;

      CategoryPath catPath = dt.getCategoryPath();
      if (catPath.getPathElements().length > 0 && categoryPathBlacklist.contains(catPath.getPathElements()[0]))
        continue;

      if (dt instanceof Structure || dt instanceof TypeDef || dt instanceof EnumDataType
          || dt instanceof Union || dt instanceof Enum || dt instanceof FunctionDefinition) {

        if (dt.getDisplayName().contains("level_displayFn") || dt.getDisplayName().contains("_M_IX86"))
          script.println("DEBUG " + dt.getDisplayName() + " - " + dt.getClass().getSimpleName());

        // if (dt.getDisplayName().contains("tdstObjectTypeElement_") ||
        // dt.getDisplayName().contains("ObjectTypeElementHandle"))

        filteredTypes.add(dt);
      }
      // }
    }

    // String s = "";
    // for (DataType dataType : filteredTypes) {
    // s += dataType.getDisplayName() + ", ";
    // }
    // script.println(s);

    try (PrintWriter writer = new PrintWriter(new File(RemanConfig.INSTANCE.outputDir, "gh_types.h"),
        "UTF-8")) {
      Utils.headerGuardPre(writer, "STRUCTS");
      writer.println("// AUTO-GENERATED FILE ");
      writer.println("#include <r3/binders/type.h>");

      DataTypeWriter dtw = new DataTypeWriter(dtm, writer);
      dtw.requiredFunctionTypes = requiredFunctionTypes;
      dtw.write(filteredTypes, script.getMonitor());

      Utils.headerGuardPost(writer, "STRUCTS");
    }
  }
}