gistfile1.java

// Use the selected value as an offset for a reference with an user-chosen base block.
// @author Alessandro Gatti
// @category References
// @keybinding
// @menupath
// @toolbar

import docking.DialogComponentProvider;
import docking.widgets.table.AbstractSortedTableModel;
import docking.widgets.table.GBooleanCellRenderer;
import docking.widgets.table.GTableCellRenderer;
import docking.widgets.table.SelectionManager;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.address.OverlayAddressSpace;
import ghidra.program.model.listing.CodeUnit;
import ghidra.program.model.listing.Data;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.listing.Program;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.program.model.mem.MemoryBlockSourceInfo;
import ghidra.program.model.mem.MemoryBlockType;
import ghidra.program.model.scalar.Scalar;
import ghidra.program.model.symbol.RefType;
import ghidra.program.model.symbol.SourceType;
import ghidra.program.util.OperandFieldLocation;
import ghidra.util.table.GhidraTable;
import ghidra.util.table.GhidraTableFilterPanel;
import java.awt.*;
import java.awt.event.KeyAdapter;
import java.awt.event.KeyEvent;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.util.*;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import javax.swing.*;
import javax.swing.table.TableModel;

public class AddAnySegmentReference extends GhidraScript {
  @Override
  protected void run() throws Exception {
    if (!(currentLocation instanceof OperandFieldLocation)) {
      return;
    }

    var dialog = new SegmentSelectionDialog(currentProgram, (OperandFieldLocation) currentLocation);
    state.getTool().showDialog(dialog);
  }

  private static class SegmentSelectionDialog extends DialogComponentProvider {

    private final Program program;

    private final SegmentsModel tableModel;

    private final GhidraTable table;

    private final OperandFieldLocation location;

    public SegmentSelectionDialog(Program program, OperandFieldLocation location) {
      super("Choose base for offset", true, true, true, false);
      this.program = program;
      this.location = location;
      tableModel = new SegmentsModel(program);
      table = buildTable(tableModel);
      addWorkPanel(buildMainPanel());
      addOKButton();
      addCancelButton();
      okButton.setText("Select");
      okButton.setEnabled(tableModel.getRowCount() > 0);
      setDefaultButton(okButton);
    }

    private JPanel buildMainPanel() {
      var panel = new JPanel(new BorderLayout());
      panel.add(new JScrollPane(table), BorderLayout.CENTER);
      panel.add(new GhidraTableFilterPanel<>(table, tableModel), BorderLayout.SOUTH);
      return panel;
    }

    private GhidraTable buildTable(SegmentsModel model) {
      var table = new SegmentsTable(model);
      table.setPreferredScrollableViewportSize(new Dimension(700, 105));
      var monoCellRenderer =
          new GTableCellRenderer() {
            @Override
            protected Font getDefaultFont() {
              return fixedWidthFont;
            }
          };
      for (var columnName :
          new String[] {
            SegmentsModel.START_COL, SegmentsModel.END_COL, SegmentsModel.LENGTH_COL,
          }) {
        table.getColumn(columnName).setCellRenderer(monoCellRenderer);
      }
      for (var columnName :
          new String[] {
            SegmentsModel.READ_COL,
            SegmentsModel.WRITE_COL,
            SegmentsModel.EXECUTE_COL,
            SegmentsModel.VOLATILE_COL,
            SegmentsModel.INIT_COL
          }) {
        table.getColumn(columnName).setCellRenderer(new GBooleanCellRenderer());
      }
      table.addKeyListener(
          new KeyAdapter() {
            @Override
            public void keyPressed(KeyEvent keyEvent) {
              if (keyEvent.getKeyCode() != KeyEvent.VK_ENTER) {
                return;
              }
              table.setRowSelectionInterval(table.getSelectedRow(), table.getSelectedRow());
              keyEvent.consume();
              okCallback();
            }
          });
      table.addMouseListener(
          new MouseAdapter() {
            @Override
            public void mouseClicked(MouseEvent mouseEvent) {
              if (mouseEvent.getClickCount() != 2) {
                return;
              }
              table.setRowSelectionInterval(table.getSelectedRow(), table.getSelectedRow());
              mouseEvent.consume();
              okCallback();
            }
          });

      var column = table.getColumn(SegmentsModel.READ_COL);
      column.setMaxWidth(25);
      column.setMinWidth(25);
      column.setResizable(false);

      column = table.getColumn(SegmentsModel.WRITE_COL);
      column.setMaxWidth(25);
      column.setMinWidth(25);
      column.setResizable(false);

      column = table.getColumn(SegmentsModel.EXECUTE_COL);
      column.setMaxWidth(25);
      column.setMinWidth(25);
      column.setResizable(false);

      column = table.getColumn(SegmentsModel.VOLATILE_COL);
      column.setMaxWidth(65);
      column.setMinWidth(65);
      column.setResizable(false);

      column = table.getColumn(SegmentsModel.BLOCK_TYPE_COL);
      column.setMinWidth(25);

      column = table.getColumn(SegmentsModel.INIT_COL);
      column.setMaxWidth(25);
      column.setMinWidth(25);
      column.setResizable(false);

      return table;
    }

    private Long getDataDisplacement(Data codeUnit) {
      var value = codeUnit.getValue();

      if (value instanceof Address address) {
        return address.getOffset();
      }

      if (value instanceof Scalar scalar) {
        return scalar.getUnsignedValue();
      }

      return null;
    }

    private Long getOperandDisplacement(CodeUnit codeUnit) {
      var operandIndex = location.getOperandIndex();

      var scalar = codeUnit.getScalar(operandIndex);
      if (scalar != null) {
        return scalar.getUnsignedValue();
      }

      var operandSubIndex = location.getSubOperandIndex();
      if (operandSubIndex >= 0) {
        var objects = ((Instruction) codeUnit).getDefaultOperandRepresentationList(operandIndex);
        if (objects != null) {
          if (objects.size() > operandSubIndex) {
            var object = objects.get(operandSubIndex);
            if (object instanceof Scalar) {
              return ((Scalar) objects.get(operandSubIndex)).getUnsignedValue();
            }
          }
          for (var object : objects) {
            if (object instanceof Scalar) {
              return ((Scalar) object).getUnsignedValue();
            }
          }
        }
      }

      return null;
    }

    @Override
    protected void okCallback() {
      var selectedRowIndex = table.getSelectedRow();
      if (selectedRowIndex < 0) {
        return;
      }
      var listing = program.getListing();
      var codeUnit = listing.getCodeUnitAt(location.getAddress());

      Long displacement = null;
      if (codeUnit instanceof Data) {
        var dataDisplacement = getDataDisplacement((Data) codeUnit);
        if (dataDisplacement == null) {
          return;
        }
        displacement = dataDisplacement;
      } else {
        if (location.getOperandIndex() >= 0) {
          var operandDisplacement = getOperandDisplacement(codeUnit);
          if (operandDisplacement == null) {
            return;
          }
          displacement = operandDisplacement;
        }
      }

      if (displacement == null) {
        return;
      }

      var referenceManager = program.getReferenceManager();
      var memoryBlock = tableModel.getRowObject(selectedRowIndex);
      var addressBase = memoryBlock.getStart();

      referenceManager.addMemoryReference(
          location.getAddress(),
          addressBase.add(displacement),
          RefType.DATA,
          SourceType.USER_DEFINED,
          location.getOperandIndex());

      close();
    }
  }

  static String getOverlayBaseSpaceName(MemoryBlock block) {
    var space = block.getStart().getAddressSpace();
    return space instanceof OverlayAddressSpace ovSpace
        ? ovSpace.getOverlayedSpace().getName()
        : "";
  }

  private static class SegmentsTable extends GhidraTable {
    SegmentsTable(TableModel model) {
      super(model);
      setVisibleRowCount(10);
      setAutoCreateColumnsFromModel(false);
    }

    @Override
    protected <T> SelectionManager createSelectionManager() {
      return null;
    }
  }

  private static class SegmentsModel extends AbstractSortedTableModel<MemoryBlock> {

    static final byte NAME = 0;
    static final byte START = 1;
    static final byte END = 2;
    static final byte LENGTH = 3;
    static final byte READ = 4;
    static final byte WRITE = 5;
    static final byte EXECUTE = 6;
    static final byte VOLATILE = 7;
    static final byte OVERLAY = 8;
    static final byte BLOCK_TYPE = 9;
    static final byte INIT = 10;
    static final byte BYTE_SOURCE = 11;
    static final byte SOURCE = 12;
    static final byte COMMENT = 13;

    static final String NAME_COL = "Name";
    static final String START_COL = "Start";
    static final String END_COL = "End";
    static final String LENGTH_COL = "Length";
    static final String READ_COL = "R";
    static final String WRITE_COL = "W";
    static final String EXECUTE_COL = "X";
    static final String VOLATILE_COL = "Volatile";
    static final String OVERLAY_COL = "Overlayed Space";
    static final String BLOCK_TYPE_COL = "Type";
    static final String INIT_COL = "Initialized";
    static final String BYTE_SOURCE_COL = "Byte Source";
    static final String SOURCE_COL = "Source";
    static final String COMMENT_COL = "Comment";

    private static final String[] COLUMN_NAMES = {
      NAME_COL,
      START_COL,
      END_COL,
      LENGTH_COL,
      READ_COL,
      WRITE_COL,
      EXECUTE_COL,
      VOLATILE_COL,
      OVERLAY_COL,
      BLOCK_TYPE_COL,
      INIT_COL,
      BYTE_SOURCE_COL,
      SOURCE_COL,
      COMMENT_COL
    };

    private final List<MemoryBlock> segmentsList;

    SegmentsModel(Program program) {
      super(START);
      segmentsList = new ArrayList<>(Arrays.stream(program.getMemory().getBlocks()).toList());
    }

    @Override
    public String getName() {
      return "Segments List";
    }

    @Override
    public List<MemoryBlock> getModelData() {
      return segmentsList;
    }

    @Override
    public String getColumnName(int column) {
      return COLUMN_NAMES[column];
    }

    @Override
    public Object getColumnValueForRow(MemoryBlock memoryBlock, int columnIndex) {
      switch (columnIndex) {
        case NAME:
          return memoryBlock.getName();

        case START:
          return memoryBlock.getStart().toString();

        case END:
          return memoryBlock.getEnd().toString();

        case LENGTH:
          return "0x" + Long.toHexString(memoryBlock.getEnd().subtract(memoryBlock.getStart()) + 1);

        case READ:
          return memoryBlock.isRead() ? Boolean.TRUE : Boolean.FALSE;

        case WRITE:
          return memoryBlock.isWrite() ? Boolean.TRUE : Boolean.FALSE;

        case EXECUTE:
          return memoryBlock.isExecute() ? Boolean.TRUE : Boolean.FALSE;

        case VOLATILE:
          return memoryBlock.isVolatile() ? Boolean.TRUE : Boolean.FALSE;

        case OVERLAY:
          return getOverlayBaseSpaceName(memoryBlock);

        case INIT:
          var blockType = memoryBlock.getType();
          return blockType == MemoryBlockType.BIT_MAPPED
              ? null
              : (memoryBlock.isInitialized() ? Boolean.TRUE : Boolean.FALSE);

        case BYTE_SOURCE:
          var sourceInfos = memoryBlock.getSourceInfos();
          var limited = sourceInfos.size() < 5 ? sourceInfos : sourceInfos.subList(0, 4);
          var description =
              limited.stream()
                  .map(MemoryBlockSourceInfo::getDescription)
                  .collect(Collectors.joining(" | "));
          if (limited != sourceInfos) {
            description += "...";
          }
          return description;

        case SOURCE:
          if ((memoryBlock.getType() == MemoryBlockType.BIT_MAPPED)
              || (memoryBlock.getType() == MemoryBlockType.BYTE_MAPPED)) {
            MemoryBlockSourceInfo info = memoryBlock.getSourceInfos().get(0);
            assert info.getMappedRange().isPresent();
            return info.getMappedRange().get().getMinAddress().toString();
          }
          return memoryBlock.getSourceName();

        case COMMENT:
          return memoryBlock.getComment();

        case BLOCK_TYPE:
          return memoryBlock.getType().toString();

        default:
          return "UNKNOWN";
      }
    }

    @Override
    public boolean isSortable(int columnIndex) {
      return columnIndex != READ
          && columnIndex != WRITE
          && columnIndex != EXECUTE
          && columnIndex != VOLATILE
          && columnIndex != INIT;
    }

    @Override
    public int getColumnCount() {
      return COLUMN_NAMES.length;
    }

    @Override
    public int findColumn(String columnName) {
      return IntStream.range(0, COLUMN_NAMES.length)
          .filter(index -> COLUMN_NAMES[index].equals(columnName))
          .findFirst()
          .orElse(0);
    }

    @Override
    protected Comparator<MemoryBlock> createSortComparator(int columnIndex) {
      return columnIndex == BYTE_SOURCE
          ? super.createSortComparator(columnIndex)
          : new MemoryMapComparator(columnIndex);
    }

    @Override
    public Class<?> getColumnClass(int columnIndex) {
      return columnIndex == READ
              || columnIndex == WRITE
              || columnIndex == EXECUTE
              || columnIndex == VOLATILE
              || columnIndex == INIT
          ? Boolean.class
          : String.class;
    }
  }

  private record MemoryMapComparator(int sortColumn) implements Comparator<MemoryBlock> {

    @Override
    public int compare(MemoryBlock first, MemoryBlock second) {
      return switch (sortColumn) {
        case SegmentsModel.NAME -> first.getName().compareToIgnoreCase(second.getName());
        case SegmentsModel.START -> first.getStart().compareTo(second.getStart());
        case SegmentsModel.END -> first.getEnd().compareTo(second.getEnd());
        case SegmentsModel.LENGTH -> (int) (first.getSize() - second.getSize());
        case SegmentsModel.READ -> ((first.isRead() ? 1 : -1) - (second.isRead() ? 1 : -1));
        case SegmentsModel.WRITE -> ((first.isWrite() ? 1 : -1) - (second.isWrite() ? 1 : -1));
        case SegmentsModel.EXECUTE ->
            ((first.isExecute() ? 1 : -1) - (second.isExecute() ? 1 : -1));
        case SegmentsModel.VOLATILE ->
            ((first.isVolatile() ? 1 : -1) - (second.isVolatile() ? 1 : -1));
        case SegmentsModel.OVERLAY ->
            getOverlayBaseSpaceName(first).compareTo(getOverlayBaseSpaceName(second));
        case SegmentsModel.INIT ->
            ((first.isInitialized() ? 1 : -1) - (second.isInitialized() ? 1 : -1));
        case SegmentsModel.SOURCE ->
            Objects.requireNonNullElse(first.getSourceName(), "")
                .compareToIgnoreCase(Objects.requireNonNullElse(second.getSourceName(), ""));
        case SegmentsModel.COMMENT ->
            Objects.requireNonNullElse(first.getComment(), "")
                .compareToIgnoreCase(Objects.requireNonNullElse(second.getComment(), ""));
        case SegmentsModel.BLOCK_TYPE ->
            first.getType().toString().compareToIgnoreCase(second.getType().toString());
        default -> throw new RuntimeException("Unimplemented column comparator: " + sortColumn);
      };
    }
  }
}