diff --git a/src/main/java/ai/reveng/toolkit/ghidra/Utils.java b/src/main/java/ai/reveng/toolkit/ghidra/Utils.java index f38c5e5b..1e23a143 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/Utils.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/Utils.java @@ -55,6 +55,43 @@ public static void addRowToDescriptor( addRowToDescriptor(descriptor, columnName, true, columnTypeClass, rowObjectAccessor); } + /** + * Helper method to add a column with sort ordinal specification. + * @param sortOrdinal 1-based sort priority (1 = primary sort), or -1 for no default sort + * @param ascending true for ascending sort, false for descending + */ + public static void addRowToDescriptor( + TableColumnDescriptor descriptor, + String columnName, + Class columnTypeClass, + RowObjectAccessor rowObjectAccessor, + int sortOrdinal, + boolean ascending) { + + var column = new AbstractDynamicTableColumn() { + @Override + public String getColumnName() { + return columnName; + } + + @Override + public COLUMN_TYPE getValue(ROW_TYPE rowObject, Settings settings, Object data, ServiceProvider serviceProvider) throws IllegalArgumentException { + return rowObjectAccessor.access(rowObject); + } + + @Override + public Class getColumnClass() { + return columnTypeClass; + } + + @Override + public Class getSupportedRowType() { + return null; + } + }; + descriptor.addVisibleColumn(column, sortOrdinal, ascending); + } + @FunctionalInterface diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java index 16e88465..11adf0e0 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/cmds/ApplyMatchCmd.java @@ -48,9 +48,7 @@ private boolean shouldApplyMatch() { return func != null && // Do not override user-defined function names func.getSymbol().getSource() != SourceType.USER_DEFINED && - // Exclude thunks and external functions - !func.isThunk() && - !func.isExternal() && + GhidraRevengService.isRelevantForAnalysis(func) && // Only accept valid names (no spaces) !match.functionMatch().nearest_neighbor_mangled_function_name().contains(" ") && !match.functionMatch().nearest_neighbor_function_name().contains(" ") diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/analysiscreation/RevEngAIAnalysisOptionsDialog.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/analysiscreation/RevEngAIAnalysisOptionsDialog.java index 30df141a..594c9b4c 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/analysiscreation/RevEngAIAnalysisOptionsDialog.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/analysiscreation/RevEngAIAnalysisOptionsDialog.java @@ -2,10 +2,12 @@ import ai.reveng.model.ConfigResponse; import ai.reveng.toolkit.ghidra.binarysimilarity.ui.dialog.RevEngDialogComponentProvider; +import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection.FunctionSelectionPanel; import ai.reveng.toolkit.ghidra.core.services.api.AnalysisOptionsBuilder; import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; import ai.reveng.toolkit.ghidra.core.services.api.types.AnalysisScope; import ai.reveng.toolkit.ghidra.plugins.ReaiPluginPackage; +import ghidra.framework.plugintool.PluginTool; import ghidra.program.model.listing.Program; import ghidra.util.Msg; import ghidra.util.Swing; @@ -25,6 +27,7 @@ public class RevEngAIAnalysisOptionsDialog extends RevEngDialogComponentProvider private JCheckBox dynamicExecutionCheckBox; private final Program program; private final GhidraRevengService service; + private final PluginTool tool; private JRadioButton privateScope; private JRadioButton publicScope; private JTextField tagsTextBox; @@ -33,22 +36,25 @@ public class RevEngAIAnalysisOptionsDialog extends RevEngDialogComponentProvider private JCheckBox identifyCVECheckBox; private JCheckBox generateSBOMCheckBox; private JComboBox architectureComboBox; + private FunctionSelectionPanel functionSelectionPanel; private boolean okPressed = false; + private boolean configCheckPassed = false; private JLabel fileSizeWarningLabel; private JLabel loadingLabel; - public static RevEngAIAnalysisOptionsDialog withModelsFromServer(Program program, GhidraRevengService reService) { - return new RevEngAIAnalysisOptionsDialog(program, reService); + public static RevEngAIAnalysisOptionsDialog withModelsFromServer(Program program, GhidraRevengService reService, PluginTool tool) { + return new RevEngAIAnalysisOptionsDialog(program, tool, reService); } - public RevEngAIAnalysisOptionsDialog(Program program, GhidraRevengService service) { + public RevEngAIAnalysisOptionsDialog(Program program, PluginTool tool, GhidraRevengService service) { super(ReaiPluginPackage.WINDOW_PREFIX + "Configure Analysis for %s".formatted(program.getName()), true); this.program = program; this.service = service; + this.tool = tool; buildInterface(); - setPreferredSize(320, 420); + setPreferredSize(600, 550); fetchConfigAsync(); } @@ -176,18 +182,26 @@ private void buildInterface() { loadingLabel.setBorder(BorderFactory.createEmptyBorder(5, 5, 5, 5)); workPanel.add(loadingLabel); + workPanel.add(new JSeparator(SwingConstants.HORIZONTAL)); + + // Add function selection panel + functionSelectionPanel = new FunctionSelectionPanel(tool); + functionSelectionPanel.initForProgram(program); + functionSelectionPanel.getTableModel().addTableModelListener(e -> updateStartButtonState()); + workPanel.add(functionSelectionPanel); + + addCancelButton(); addOKButton(); okButton.setText("Start Analysis"); - okButton.setEnabled(false); // Disabled until config check completes + okButton.setEnabled(false); // Disabled until config check completes and functions are selected } - public @Nullable AnalysisOptionsBuilder getOptionsFromUI() { - if (!okPressed) { - return null; - } - var options = AnalysisOptionsBuilder.forProgram(program); + public AnalysisOptionsBuilder getOptionsFromUI() { + // Use the selected functions from the function selection panel + var selectedFunctions = functionSelectionPanel.getSelectedFunctions(); + var options = AnalysisOptionsBuilder.forProgramWithFunctions(program, selectedFunctions); options.skipScraping(!scrapeExternalTagsBox.isSelected()); options.skipCapabilities(!identifyCapabilitiesCheckBox.isSelected()); @@ -216,6 +230,10 @@ protected void okCallback() { close(); } + public boolean isOkPressed() { + return okPressed; + } + @Override public JComponent getComponent() { return super.getComponent(); @@ -239,7 +257,8 @@ private void handleConfigResponse(@Nullable ConfigResponse config) { if (config == null) { // Config fetch failed, allow upload attempt (server will reject if too large) - okButton.setEnabled(true); + configCheckPassed = true; + updateStartButtonState(); return; } @@ -251,7 +270,8 @@ private void validateFileSize(long maxFileSizeBytes) { long fileSize = getProgramFileSize(); if (fileSize < 0) { // Could not determine file size, allow upload attempt - okButton.setEnabled(true); + configCheckPassed = true; + updateStartButtonState(); return; } @@ -262,13 +282,19 @@ private void validateFileSize(long maxFileSizeBytes) { "
File size (%s) exceeds
server limit (%s)
" .formatted(fileSizeStr, maxSizeStr)); fileSizeWarningLabel.setVisible(true); - okButton.setEnabled(false); + configCheckPassed = false; + updateStartButtonState(); } else { fileSizeWarningLabel.setVisible(false); - okButton.setEnabled(true); + configCheckPassed = true; + updateStartButtonState(); } } + private void updateStartButtonState() { + okButton.setEnabled(configCheckPassed && functionSelectionPanel.getSelectedCount() > 0); + } + private long getProgramFileSize() { try { Path filePath; diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionRowObject.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionRowObject.java new file mode 100644 index 00000000..6fc45093 --- /dev/null +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionRowObject.java @@ -0,0 +1,111 @@ +package ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection; + +import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionInfo; +import ghidra.program.model.address.Address; +import ghidra.program.model.listing.Function; + +import javax.annotation.Nullable; + +/** + * Wrapper around a Ghidra {@link Function} with a mutable selection flag. + * Used to display functions in a table where users can select which functions + * to include in analysis. + */ +public class FunctionRowObject { + private final Function function; + private boolean selected; + private boolean enabled = true; + @Nullable + private FunctionInfo remoteFunctionInfo; + + public FunctionRowObject(Function function, boolean selected) { + this.function = function; + this.selected = selected; + } + + public Function getFunction() { + return function; + } + + public String getName() { + return function.getName(); + } + + public Address getAddress() { + return function.getEntryPoint(); + } + + /** + * Returns the size of the function based on address count. + */ + public long getSize() { + return function.getBody().getNumAddresses(); + } + + public boolean isExternal() { + return function.isExternal(); + } + + public boolean isThunk() { + return function.isThunk(); + } + + /** + * Returns a human-readable type string for the function. + */ + public String getType() { + if (isExternal()) { + return "External"; + } else if (isThunk()) { + return "Thunk"; + } else { + return "Normal"; + } + } + + public boolean isSelected() { + return selected; + } + + public void setSelected(boolean selected) { + if (!enabled) { + return; + } + this.selected = selected; + } + + public boolean isEnabled() { + return enabled; + } + + public void setEnabled(boolean enabled) { + this.enabled = enabled; + if (!enabled) { + this.selected = false; + } + } + + @Nullable + public FunctionInfo getRemoteFunctionInfo() { + return remoteFunctionInfo; + } + + public void setRemoteFunctionInfo(@Nullable FunctionInfo remoteFunctionInfo) { + this.remoteFunctionInfo = remoteFunctionInfo; + } + + @Nullable + public String getRemoteFunctionName() { + return remoteFunctionInfo != null ? remoteFunctionInfo.functionName() : null; + } + + @Nullable + public String getRemoteMangledName() { + return remoteFunctionInfo != null ? remoteFunctionInfo.functionMangledName() : null; + } + + @Nullable + public Long getRemoteFunctionID() { + return remoteFunctionInfo != null ? remoteFunctionInfo.functionID().value() : null; + } +} diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionSelectionPanel.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionSelectionPanel.java new file mode 100644 index 00000000..b058cccd --- /dev/null +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionSelectionPanel.java @@ -0,0 +1,186 @@ +package ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection; + +import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionInfo; +import ghidra.framework.plugintool.ServiceProvider; +import ghidra.program.model.listing.Function; +import ghidra.program.model.listing.Program; +import ghidra.util.table.GhidraFilterTable; + +import javax.swing.*; +import javax.swing.border.TitledBorder; +import javax.swing.event.TableModelEvent; +import javax.swing.event.TableModelListener; +import java.awt.*; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * A reusable panel for selecting functions from a Ghidra program. + * Contains a filterable table of functions with selection checkboxes, + * toolbar buttons for bulk selection operations, and a summary label. + */ +public class FunctionSelectionPanel extends JPanel { + private final FunctionSelectionTableModel tableModel; + private final GhidraFilterTable filterTable; + private final JLabel summaryLabel; + + public FunctionSelectionPanel(ServiceProvider serviceProvider) { + super(new BorderLayout()); + + tableModel = new FunctionSelectionTableModel(serviceProvider); + filterTable = new GhidraFilterTable<>(tableModel); + summaryLabel = new JLabel(); + + buildInterface(); + + // Listen for table changes to update the summary + tableModel.addTableModelListener(new TableModelListener() { + @Override + public void tableChanged(TableModelEvent e) { + updateSummaryLabel(); + } + }); + } + + private void buildInterface() { + setBorder(new TitledBorder("Function Selection")); + + // Toolbar with bulk selection buttons + JPanel toolbarPanel = new JPanel(new FlowLayout(FlowLayout.LEFT)); + + JButton selectAllButton = new JButton("Select All"); + selectAllButton.setName("selectAllButton"); + selectAllButton.addActionListener(e -> { + tableModel.selectAll(); + updateSummaryLabel(); + }); + + JButton deselectAllButton = new JButton("Deselect All"); + deselectAllButton.setName("deselectAllButton"); + deselectAllButton.addActionListener(e -> { + tableModel.deselectAll(); + updateSummaryLabel(); + }); + + JButton excludeUserDefinedButton = new JButton("Exclude User-Defined"); + excludeUserDefinedButton.setName("excludeUserDefinedButton"); + excludeUserDefinedButton.setToolTipText("Deselect functions with user-defined name or signature"); + excludeUserDefinedButton.addActionListener(e -> { + tableModel.deselectUserDefined(); + updateSummaryLabel(); + }); + + toolbarPanel.add(selectAllButton); + toolbarPanel.add(deselectAllButton); + toolbarPanel.add(excludeUserDefinedButton); + + // Summary label on the right side of the toolbar + JPanel summaryPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT)); + summaryPanel.add(summaryLabel); + + JPanel topPanel = new JPanel(new BorderLayout()); + topPanel.add(toolbarPanel, BorderLayout.WEST); + topPanel.add(summaryPanel, BorderLayout.EAST); + + add(topPanel, BorderLayout.NORTH); + add(filterTable, BorderLayout.CENTER); + + updateSummaryLabel(); + } + + /** + * Initialize the panel with functions from the given program. + * External and thunk functions are excluded. By default, all functions are selected. + */ + public void initForProgram(Program program) { + tableModel.initForProgram(program); + updateSummaryLabel(); + } + + /** + * Returns the list of currently selected functions. + */ + public List getSelectedFunctions() { + return tableModel.getSelectedFunctions(); + } + + /** + * Returns the count of selected functions. + */ + public int getSelectedCount() { + return tableModel.getSelectedCount(); + } + + /** + * Returns the total number of functions. + */ + public int getTotalFunctionCount() { + return tableModel.getTotalCount(); + } + + /** + * Returns the underlying table model. + */ + public FunctionSelectionTableModel getTableModel() { + return tableModel; + } + + /** + * Returns the underlying filter table component. + */ + public GhidraFilterTable getFilterTable() { + return filterTable; + } + + /** + * Cross-reference local functions with remote function info. + * Functions without a remote match (or with a size mismatch) are disabled. + */ + public void applyRemoteFunctionInfo(List remoteFunctions) { + // Build lookup by virtual address + Map byAddress = new HashMap<>(); + for (FunctionInfo info : remoteFunctions) { + byAddress.put(info.functionVirtualAddress(), info); + } + + for (FunctionRowObject row : tableModel.getAllRows()) { + // Reset state from any previous matching + row.setEnabled(true); + row.setRemoteFunctionInfo(null); + + long localAddr = row.getAddress().getOffset(); + FunctionInfo match = byAddress.get(localAddr); + if (match != null && sizeMatches(row, match)) { + row.setRemoteFunctionInfo(match); + row.setSelected(true); + } else { + row.setEnabled(false); + } + } + + tableModel.fireTableDataChanged(); + updateSummaryLabel(); + } + + /** + * Check if the local function size matches the remote function size, + * using the same off-by-one tolerance as GhidraRevengService. + */ + private static boolean sizeMatches(FunctionRowObject row, FunctionInfo info) { + long localSize = row.getSize(); + int remoteSize = info.functionSize(); + return localSize == remoteSize || localSize - 1 == remoteSize; + } + + private void updateSummaryLabel() { + int selected = tableModel.getSelectedCount(); + int total = tableModel.getTotalCount(); + int enabled = tableModel.getEnabledCount(); + if (enabled < total) { + summaryLabel.setText(String.format("%d of %d functions selected (%d matched remotely)", selected, total, enabled)); + } else { + summaryLabel.setText(String.format("%d of %d functions selected", selected, total)); + } + } +} diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionSelectionTableModel.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionSelectionTableModel.java new file mode 100644 index 00000000..ac0fbb7d --- /dev/null +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/functionselection/FunctionSelectionTableModel.java @@ -0,0 +1,211 @@ +package ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection; + +import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; +import docking.widgets.table.TableColumnDescriptor; +import docking.widgets.table.threaded.ThreadedTableModelStub; +import ghidra.framework.plugintool.ServiceProvider; +import ghidra.program.model.address.Address; +import ghidra.program.model.listing.Function; +import ghidra.program.model.listing.FunctionSignature; +import ghidra.program.model.listing.Program; +import ghidra.program.model.symbol.Namespace; +import ghidra.program.model.symbol.SourceType; +import ghidra.util.datastruct.Accumulator; +import ghidra.util.task.TaskMonitor; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static ai.reveng.toolkit.ghidra.Utils.addRowToDescriptor; + +/** + * Table model for displaying and selecting functions from a Ghidra program. + * The model allows users to select which functions should be included in analysis. + */ +public class FunctionSelectionTableModel extends ThreadedTableModelStub { + // Column index for editable Select column + static final int SELECT = 0; + + private final List functionList = new ArrayList<>(); + + public FunctionSelectionTableModel(ServiceProvider serviceProvider) { + super("Function Selection Table Model", serviceProvider); + } + + /** + * Initialize the model with functions from the given program. + * External and thunk functions are excluded from the list entirely. + * By default, all functions are selected. + */ + public void initForProgram(Program program) { + functionList.clear(); + + if (program != null) { + program.getFunctionManager().getFunctions(true).forEach(function -> { + if (!GhidraRevengService.isRelevantForAnalysis(function)) { + return; + } + functionList.add(new FunctionRowObject(function, true)); + }); + } + reload(); + } + + @Override + protected void doLoad(Accumulator accumulator, TaskMonitor monitor) { + monitor.setMessage("Loading functions"); + monitor.setMaximum(functionList.size()); + int count = 0; + for (FunctionRowObject row : functionList) { + if (monitor.isCancelled()) { + break; + } + accumulator.add(row); + monitor.setProgress(++count); + } + } + + @Override + protected TableColumnDescriptor createTableColumnDescriptor() { + TableColumnDescriptor descriptor = new TableColumnDescriptor<>(); + addRowToDescriptor(descriptor, "Select", Boolean.class, FunctionRowObject::isSelected); + addRowToDescriptor(descriptor, "Address", Address.class, FunctionRowObject::getAddress, 1, true); + addRowToDescriptor(descriptor, "Namespace", false, Namespace.class, fo -> fo.getFunction().getParentNamespace()); + addRowToDescriptor(descriptor, "Name", String.class, FunctionRowObject::getName); + addRowToDescriptor(descriptor, "Name Source", SourceType.class, fo -> fo.getFunction().getSymbol().getSource()); + addRowToDescriptor(descriptor, "Signature", false, FunctionSignature.class, fo -> fo.getFunction().getSignature()); + addRowToDescriptor(descriptor, "Signature Source", SourceType.class, fo -> fo.getFunction().getSignatureSource()); + addRowToDescriptor(descriptor, "Size", false, Long.class, FunctionRowObject::getSize); + addRowToDescriptor(descriptor, "Remote Name", String.class, FunctionRowObject::getRemoteFunctionName); + addRowToDescriptor(descriptor, "Remote Mangled Name", false, String.class, FunctionRowObject::getRemoteMangledName); + addRowToDescriptor(descriptor, "Remote Function ID", false, Long.class, FunctionRowObject::getRemoteFunctionID); + return descriptor; + } + + @Override + public boolean isCellEditable(int rowIndex, int columnIndex) { + if (columnIndex != SELECT) { + return false; + } + FunctionRowObject row = getRowObject(rowIndex); + return row.isEnabled(); + } + + @Override + public void setValueAt(Object aValue, int rowIndex, int columnIndex) { + if (columnIndex == SELECT && aValue instanceof Boolean) { + FunctionRowObject row = getRowObject(rowIndex); + if (!row.isEnabled()) { + return; + } + row.setSelected((Boolean) aValue); + fireTableRowsUpdated(rowIndex, rowIndex); + } + } + + /** + * Select all enabled functions in the table. + */ + public void selectAll() { + for (FunctionRowObject row : functionList) { + if (row.isEnabled()) { + row.setSelected(true); + } + } + fireTableDataChanged(); + } + + /** + * Deselect all functions in the table. + */ + public void deselectAll() { + for (FunctionRowObject row : functionList) { + row.setSelected(false); + } + fireTableDataChanged(); + } + + /** + * Select all non-thunk enabled functions. + * Thunk and disabled functions will be deselected. + * (External functions are not included in the list.) + */ + public void selectNonThunk() { + for (FunctionRowObject row : functionList) { + if (row.isEnabled()) { + row.setSelected(!row.isThunk()); + } + } + fireTableDataChanged(); + } + + /** + * Deselect functions that have user-defined name or signature source. + * These are functions where the user has already made manual changes. + */ + public void deselectUserDefined() { + for (FunctionRowObject row : functionList) { + var func = row.getFunction(); + var nameSource = func.getSymbol().getSource(); + var sigSource = func.getSignatureSource(); + if (nameSource == SourceType.USER_DEFINED || sigSource == SourceType.USER_DEFINED) { + row.setSelected(false); + } + } + fireTableDataChanged(); + } + + /** + * Returns the list of selected functions. + */ + public List getSelectedFunctions() { + List selected = new ArrayList<>(); + for (FunctionRowObject row : functionList) { + if (row.isSelected()) { + selected.add(row.getFunction()); + } + } + return selected; + } + + /** + * Returns the count of selected functions. + */ + public int getSelectedCount() { + int count = 0; + for (FunctionRowObject row : functionList) { + if (row.isSelected()) { + count++; + } + } + return count; + } + + /** + * Returns the total number of functions. + */ + public int getTotalCount() { + return functionList.size(); + } + + /** + * Returns an unmodifiable view of all rows. + */ + public List getAllRows() { + return Collections.unmodifiableList(functionList); + } + + /** + * Returns the count of enabled (matched remotely) rows. + */ + public int getEnabledCount() { + int count = 0; + for (FunctionRowObject row : functionList) { + if (row.isEnabled()) { + count++; + } + } + return count; + } +} diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysesTableModel.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysesTableModel.java index c3452d9a..3683463c 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysesTableModel.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysesTableModel.java @@ -3,8 +3,10 @@ import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface; import ai.reveng.toolkit.ghidra.core.services.api.types.*; -import ai.reveng.toolkit.ghidra.core.services.function.export.ExportFunctionBoundariesService; import ai.reveng.toolkit.ghidra.core.services.logging.ReaiLoggingService; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import docking.widgets.table.AbstractDynamicTableColumn; import docking.widgets.table.TableColumnDescriptor; import docking.widgets.table.threaded.ThreadedTableModelStub; @@ -19,6 +21,7 @@ public class RecentAnalysesTableModel extends ThreadedTableModelStub { private final TypedApiInterface.BinaryHash hash; private final Address imageBase; + private final Map functionCountCache = new ConcurrentHashMap<>(); public RecentAnalysesTableModel(PluginTool tool, TypedApiInterface.BinaryHash hash, Address imageBase) { super("Recent Analyses Table Model", tool); @@ -29,7 +32,6 @@ public RecentAnalysesTableModel(PluginTool tool, TypedApiInterface.BinaryHash ha @Override protected void doLoad(Accumulator accumulator, TaskMonitor monitor) throws CancelledException { var revEngAIService = serviceProvider.getService(GhidraRevengService.class); - var functionBoundariesService = serviceProvider.getService(ExportFunctionBoundariesService.class); var loggingService = serviceProvider.getService(ReaiLoggingService.class); // The search endpoint only returns analyses we have access to so there is no need to filter them. @@ -49,11 +51,26 @@ protected void doLoad(Accumulator accumulator, TaskMonitor return; } + // Fetch the function count while we're on the background thread + try { + var functions = revEngAIService.getApi().getFunctionInfo(result.analysis_id()); + functionCountCache.put(result.analysis_id(), functions.size()); + } catch (Exception e) { + loggingService.info("[RevEng] Could not fetch function count for " + result.analysis_id() + ": " + e.getMessage()); + } + accumulator.add(result); } ); } + /** + * Returns the cached remote function count for the given analysis, or null if not yet loaded. + */ + public Integer getFunctionCount(TypedApiInterface.AnalysisID analysisID) { + return functionCountCache.get(analysisID); + } + @Override protected TableColumnDescriptor createTableColumnDescriptor() { TableColumnDescriptor descriptor = new TableColumnDescriptor<>(); @@ -109,6 +126,17 @@ public AnalysisStatus getValue(LegacyAnalysisResult rowObject, Settings settings } }); + descriptor.addVisibleColumn(new AbstractDynamicTableColumn() { + @Override + public String getColumnName() { + return "Functions"; + } + + @Override + public Integer getValue(LegacyAnalysisResult rowObject, Settings settings, Object data, ServiceProvider serviceProvider) throws IllegalArgumentException { + return functionCountCache.get(rowObject.analysis_id()); + } + }); return descriptor; } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysisDialog.java b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysisDialog.java index 51f1d487..77c7592d 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysisDialog.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/binarysimilarity/ui/recentanalyses/RecentAnalysisDialog.java @@ -1,20 +1,25 @@ package ai.reveng.toolkit.ghidra.binarysimilarity.ui.recentanalyses; import ai.reveng.toolkit.ghidra.binarysimilarity.ui.dialog.RevEngDialogComponentProvider; -import ai.reveng.toolkit.ghidra.core.RevEngAIAnalysisStatusChangedEvent; +import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection.FunctionSelectionPanel; import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface; import ai.reveng.toolkit.ghidra.core.services.api.types.LegacyAnalysisResult; +import ai.reveng.toolkit.ghidra.core.tasks.AttachToAnalysisTask; import ai.reveng.toolkit.ghidra.plugins.ReaiPluginPackage; import ghidra.framework.plugintool.PluginTool; import ghidra.program.model.listing.Program; +import ghidra.util.Msg; +import ghidra.util.Swing; import ghidra.util.table.GhidraFilterTable; +import ghidra.util.task.TaskBuilder; import javax.swing.*; import java.awt.*; import java.awt.event.MouseAdapter; import java.awt.event.MouseEvent; import java.util.Comparator; +import java.util.concurrent.CompletableFuture; /** @@ -27,6 +32,10 @@ public class RecentAnalysisDialog extends RevEngDialogComponentProvider { private final PluginTool tool; private final Program program; private final GhidraRevengService ghidraRevengService; + private final FunctionSelectionPanel functionSelectionPanel; + private JButton pickMostRecentButton; + private JButton pickSelectedButton; + private volatile TypedApiInterface.AnalysisID lastFetchedAnalysisID = null; public RecentAnalysisDialog(PluginTool tool, Program program) { super(ReaiPluginPackage.WINDOW_PREFIX + "Recent Analyses", true); @@ -38,8 +47,18 @@ public RecentAnalysisDialog(PluginTool tool, Program program) { recentAnalysesTableModel = new RecentAnalysesTableModel(tool, hash, this.program.getImageBase()); recentAnalysesTable = new GhidraFilterTable<>(recentAnalysesTableModel); + functionSelectionPanel = new FunctionSelectionPanel(tool); + functionSelectionPanel.initForProgram(program); + buildInterface(); - setPreferredSize(600, 400); + setPreferredSize(1000, 700); + + // When the analyses table finishes loading, auto-select the most recent and fetch its remote functions + recentAnalysesTableModel.addTableModelListener(e -> { + if (lastFetchedAnalysisID == null && recentAnalysesTableModel.getRowCount() > 0) { + selectMostRecentRow(); + } + }); } private void buildInterface() { @@ -49,7 +68,8 @@ private void buildInterface() { JPanel titlePanel = createTitlePanel("Find existing analyses for this binary"); mainPanel.add(titlePanel, BorderLayout.NORTH); - // Create the table content + // Create the analysis table panel + JPanel analysisTablePanel = new JPanel(new BorderLayout()); // Add mouse listener to handle clicks on the Analysis ID column recentAnalysesTable.getTable().addMouseListener(new MouseAdapter() { @Override @@ -72,9 +92,26 @@ public void mouseClicked(MouseEvent e) { } } }); - mainPanel.add(recentAnalysesTable, BorderLayout.CENTER); + // When the user selects a different analysis row, fetch its remote functions + recentAnalysesTable.getTable().getSelectionModel().addListSelectionListener(e -> { + if (e.getValueIsAdjusting()) { + return; + } + var selected = recentAnalysesTable.getSelectedRowObject(); + if (selected != null) { + fetchRemoteFunctions(selected.analysis_id()); + } + }); + + analysisTablePanel.add(recentAnalysesTable, BorderLayout.CENTER); + + // Create split pane with analysis table on top and function selection on bottom + JSplitPane splitPane = new JSplitPane(JSplitPane.VERTICAL_SPLIT, analysisTablePanel, functionSelectionPanel); + splitPane.setResizeWeight(0.4); // Give 40% to analysis table, 60% to function selection + splitPane.setDividerLocation(200); + mainPanel.add(splitPane, BorderLayout.CENTER); - JButton pickMostRecentButton = new JButton("Pick most recent"); + pickMostRecentButton = new JButton("Pick most recent"); pickMostRecentButton.setName("Pick most recent"); pickMostRecentButton.addActionListener(e -> { var mostRecent = recentAnalysesTable.getModel().getModelData().stream().max( @@ -84,7 +121,7 @@ public void mouseClicked(MouseEvent e) { }); addButton(pickMostRecentButton); - JButton pickSelectedButton = new JButton("Pick selected"); + pickSelectedButton = new JButton("Pick selected"); pickSelectedButton.setName("Pick selected"); pickSelectedButton.addActionListener(e -> { var selectedRowObject = recentAnalysesTable.getSelectedRowObject(); @@ -92,20 +129,74 @@ public void mouseClicked(MouseEvent e) { }); addButton(pickSelectedButton); + functionSelectionPanel.getTableModel().addTableModelListener(e -> updatePickButtonsState()); + updatePickButtonsState(); + addWorkPanel(mainPanel); } + private void updatePickButtonsState() { + boolean hasSelection = functionSelectionPanel.getSelectedCount() > 0; + pickMostRecentButton.setEnabled(hasSelection); + pickSelectedButton.setEnabled(hasSelection); + } + + private void selectMostRecentRow() { + var modelData = recentAnalysesTableModel.getModelData(); + var mostRecent = modelData.stream() + .max(Comparator.comparing(LegacyAnalysisResult::creation)) + .orElse(null); + if (mostRecent == null) { + return; + } + // Find the view index for this row and select it — triggers the selection listener + var table = recentAnalysesTable.getTable(); + for (int i = 0; i < table.getRowCount(); i++) { + if (recentAnalysesTable.getModel().getRowObject(i) == mostRecent) { + table.setRowSelectionInterval(i, i); + break; + } + } + } + + private void fetchRemoteFunctions(TypedApiInterface.AnalysisID analysisID) { + if (analysisID.equals(lastFetchedAnalysisID)) { + return; + } + lastFetchedAnalysisID = analysisID; + var api = ghidraRevengService.getApi(); + + CompletableFuture.supplyAsync(() -> api.getFunctionInfo(analysisID)) + .thenAccept(remoteFunctions -> Swing.runLater(() -> { + functionSelectionPanel.applyRemoteFunctionInfo(remoteFunctions); + updatePickButtonsState(); + })) + .exceptionally(ex -> { + Msg.warn(this, "Failed to fetch remote function info for " + analysisID + ": " + ex.getMessage()); + return null; + }); + } + + /// Currently [[ai.reveng.toolkit.ghidra.binarysimilarity.ui.recentanalyses.RecentAnalysesTableModel#doLoad]] + /// only allows selecting a complete analysis. This simplifies the logic around the function selection panel. + /// private void pickAnalysis(LegacyAnalysisResult result) { var service = tool.getService(GhidraRevengService.class); var analysisID = service.getApi().getAnalysisIDfromBinaryID(result.binary_id()); - var programWithID = service.registerAnalysisForProgram(program, analysisID); - tool.firePluginEvent( - new RevEngAIAnalysisStatusChangedEvent( - "Recent Analysis Dialog", - programWithID, - result.status() - ) - ); + // Register the analysis ID with the program (persists to program options) + var programWithId = service.registerAnalysisForProgram(program, analysisID); + + // Get the selected functions from the function selection panel + var selectedFunctions = functionSelectionPanel.getSelectedFunctions(); + + // Create and run the attach task modally - blocks until complete + var task = new AttachToAnalysisTask(programWithId, selectedFunctions, service, tool); + TaskBuilder.withTask(task) + .setCanCancel(false) + .setStatusTextAlignment(SwingConstants.LEADING) + .launchModal(); + + // Close the dialog after task completes close(); } } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/AnalysisOptionsBuilder.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/AnalysisOptionsBuilder.java index ad720a97..0d9f5bc9 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/AnalysisOptionsBuilder.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/AnalysisOptionsBuilder.java @@ -2,6 +2,7 @@ import ai.reveng.toolkit.ghidra.core.services.api.types.AnalysisScope; import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionBoundary; +import ghidra.program.model.listing.Function; import ghidra.program.model.listing.Program; import ghidra.util.Msg; import org.json.JSONArray; @@ -34,6 +35,26 @@ public AnalysisOptionsBuilder functionBoundaries(long base, List selectedFunctions) { + List boundaries = selectedFunctions.stream() + .filter(GhidraRevengService::isRelevantForAnalysis) + .map(function -> new FunctionBoundary( + function.getSymbol().getName(false), + function.getEntryPoint().getOffset(), + function.getBody().getMaxAddress().getOffset() + )) + .toList(); + return functionBoundaries(base, boundaries); + } + public AnalysisOptionsBuilder hash(TypedApiInterface.BinaryHash hash) { options.put("sha_256_hash", hash.sha256()); return this; @@ -74,6 +95,24 @@ public static AnalysisOptionsBuilder forProgram(Program program) { ); } + /** + * Creates an AnalysisOptionsBuilder for a program with a specific list of functions. + * Only the specified functions will be included in the analysis. + * + * @param program The Ghidra program + * @param selectedFunctions The list of functions to include in the analysis + * @return A new AnalysisOptionsBuilder configured for the program with filtered functions + */ + public static AnalysisOptionsBuilder forProgramWithFunctions(Program program, List selectedFunctions) { + return new AnalysisOptionsBuilder() + .hash(new TypedApiInterface.BinaryHash(program.getExecutableSHA256())) + .fileName(program.getName()) + .functionBoundariesFromGhidraFunctions( + program.getImageBase().getOffset(), + selectedFunctions + ); + } + public AnalysisOptionsBuilder skipSBOM(boolean b) { options.put("skip_sbom", b); return this; diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java index 41aea3f0..1fe60feb 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/services/api/GhidraRevengService.java @@ -108,7 +108,35 @@ public ProgramWithID registerAnalysisForProgram(Program program, TypedApiInterfa return addAnalysisIDtoProgramOptions(program, analysisID); } + /** + * Registers an analysis for a program and stores a function filter for later use. + * The filter will be applied when registerFinishedAnalysisForProgram is called. + * + * @param program The program to register + * @param analysisID The analysis ID to associate + * @param selectedFunctions The functions to include when mapping (null for all functions) + * @return The program with associated analysis ID + */ + public ProgramWithID registerAnalysisForProgram(Program program, TypedApiInterface.AnalysisID analysisID, + @Nullable List selectedFunctions) { + return addAnalysisIDtoProgramOptions(program, analysisID); + } + public AnalysedProgram registerFinishedAnalysisForProgram(ProgramWithID programWithID, TaskMonitor monitor) throws CancelledException { + // Check if there's a pending function filter for this analysis + return registerFinishedAnalysisForProgram(programWithID, null, monitor); + } + + /** + * Registers a finished analysis for a program, optionally filtering which functions get mapped. + * + * @param programWithID The program with associated analysis ID + * @param selectedFunctions Optional list of functions to include. If null, all functions are included. + * @param monitor Task monitor for cancellation + * @return The analysed program with function ID mappings + */ + public AnalysedProgram registerFinishedAnalysisForProgram(ProgramWithID programWithID, + @Nullable List selectedFunctions, TaskMonitor monitor) throws CancelledException { var status = status(programWithID); if (!status.equals(AnalysisStatus.Complete)){ throw new IllegalStateException("Analysis %s is not complete yet, current status: %s" @@ -116,7 +144,15 @@ public AnalysedProgram registerFinishedAnalysisForProgram(ProgramWithID programW } statusCache.put(programWithID.analysisID, AnalysisStatus.Complete); - var analysedProgram = associateFunctionInfo(programWithID); + // Convert selected functions to a set of entry point addresses for filtering + Set
functionFilter = null; + if (selectedFunctions != null) { + functionFilter = selectedFunctions.stream() + .map(Function::getEntryPoint) + .collect(Collectors.toSet()); + } + + var analysedProgram = associateFunctionInfo(programWithID, functionFilter, monitor); pullFunctionInfoFromAnalysis(analysedProgram, monitor); monitor.checkCancelled(); return analysedProgram; @@ -255,11 +291,13 @@ private Optional getBinaryIDfromOptions( /// analysis is associated with the program /// Other function information like the name and signature should be loaded in [#pullFunctionInfoFromAnalysis(AnalysedProgram ,TaskMonitor)] /// because this information can change on the server, and thus needs a dedicated method to refresh it - private AnalysedProgram associateFunctionInfo(ProgramWithID knownProgram) { + private AnalysedProgram associateFunctionInfo(ProgramWithID knownProgram, @Nullable Set
functionFilter, TaskMonitor monitor) throws CancelledException { var analysisID = knownProgram.analysisID(); var program = knownProgram.program(); List functionInfo = null; functionInfo = api.getFunctionInfo(analysisID); + + monitor.checkCancelled(); var transactionID = program.startTransaction("Associate Function Info"); // Create the FunctionID map @@ -282,6 +320,7 @@ private AnalysedProgram associateFunctionInfo(ProgramWithID knownProgram) { LongPropertyMap finalFunctionIDMap = functionIDMap; int ghidraBoundariesMatchedFunction = 0; + int skippedByFilter = 0; for (FunctionInfo info : functionInfo) { var oFunc = getFunctionFor(info, program); if (oFunc.isEmpty()) { @@ -289,6 +328,13 @@ private AnalysedProgram associateFunctionInfo(ProgramWithID knownProgram) { continue; } var func = oFunc.get(); + + // Skip functions not in the filter (if filter is provided) + if (functionFilter != null && !functionFilter.contains(func.getEntryPoint())) { + skippedByFilter++; + continue; + } + // There are two ways to think about the size of a function // They diverge for non-contiguous functions var funcSizeByAddressCount = func.getBody().getNumAddresses(); @@ -313,7 +359,7 @@ private AnalysedProgram associateFunctionInfo(ProgramWithID knownProgram) { AtomicInteger ghidraFunctionCount = new AtomicInteger(); program.getFunctionManager().getFunctions(true).forEach( func -> { - if (!func.isExternal() && !func.isThunk()){ + if (isRelevantForAnalysis(func)){ ghidraFunctionCount.getAndIncrement(); if (analysedProgram.getIDForFunction(func).isEmpty()) { @@ -323,12 +369,14 @@ private AnalysedProgram associateFunctionInfo(ProgramWithID knownProgram) { } ); // Print summary + String filterInfo = functionFilter != null ? " (%d skipped by filter)".formatted(skippedByFilter) : ""; Msg.showInfo(this, null, ReaiPluginPackage.WINDOW_PREFIX + "Function loading summary", ("Found %d functions from RevEng.AI. Your local Ghidra instance has %d/%d matching function " + - "boundaries. For better results, please start a new analysis from this plugin.").formatted( + "boundaries%s. For better results, please start a new analysis from this plugin.").formatted( functionInfo.size(), ghidraBoundariesMatchedFunction, - ghidraFunctionCount.get() + ghidraFunctionCount.get(), + filterInfo )); return analysedProgram; @@ -348,7 +396,7 @@ public String virtualAddress() { /// * the type signature of the function /// /// It assumes that the initial load already happened, i.e. the functions have an associated FunctionID already. - /// The initial association happens in {@link #associateFunctionInfo(ProgramWithID)} + /// The initial association happens in {@link #associateFunctionInfo(ProgramWithID, Set, TaskMonitor)} /// public List pullFunctionInfoFromAnalysis(AnalysedProgram analysedProgram, TaskMonitor monitor) { var transactionId = analysedProgram.program().startTransaction("RevEng.AI: Pull Function Info from Analysis"); @@ -384,8 +432,7 @@ public List pullFunctionInfoFromAnalysis(AnalysedProgram analysedP continue; } var ghidraMangledName = function.getSymbol().getName(false); - // Skip external and thunk functions because we don't support them - if (function.isExternal() || function.isThunk()) { + if (!isRelevantForAnalysis(function)) { Msg.debug(this, "Skipping external/thunk function %s".formatted(ghidraMangledName)); continue; } @@ -559,11 +606,22 @@ private boolean isProgramAnalysed(Program program){ program.getUsrPropertyManager().getStringPropertyMap(REAI_FUNCTION_MANGLED_MAP) != null; } + /** + * Returns whether a function is relevant for sending to the RevEng.AI backend. + * External and thunk functions are excluded because the backend cannot process them. + */ + public static boolean isRelevantForAnalysis(Function function) { + return !function.isExternal() && !function.isThunk(); + } + public static List exportFunctionBoundaries(Program program){ List result = new ArrayList<>(); Address imageBase = program.getImageBase(); program.getFunctionManager().getFunctions(true).forEach( function -> { + if (!isRelevantForAnalysis(function)) { + return; + } var start = function.getEntryPoint(); var end = function.getBody().getMaxAddress(); result.add(new FunctionBoundary(function.getSymbol().getName(false), start.getOffset(), end.getOffset())); @@ -572,6 +630,23 @@ public static List exportFunctionBoundaries(Program program){ return result; } + /** + * Export function boundaries for a specific list of functions. + * + * @param program The program containing the functions + * @param functions The list of functions to export + * @return List of function boundaries for the specified functions + */ + public static List exportFunctionBoundaries(Program program, List functions) { + List result = new ArrayList<>(); + for (Function function : functions) { + var start = function.getEntryPoint(); + var end = function.getBody().getMaxAddress(); + result.add(new FunctionBoundary(function.getSymbol().getName(false), start.getOffset(), end.getOffset())); + } + return result; + } + private TypedApiInterface.BinaryHash hashOfProgram(Program program) { // TODO: we break the guarantee that a BinaryHash implies that a file of this hash has already been uploaded return new TypedApiInterface.BinaryHash(program.getExecutableSHA256()); @@ -942,6 +1017,10 @@ public void openPortalFor(TypedApiInterface.FunctionID f){ openFunctionInPortal(f); } + public void openPortalFor(FunctionWithID functionWithID) { + openPortalFor(functionWithID.functionID); + } + public void openPortalFor(AnalysisResult analysisResult) { openPortalFor(analysisResult.analysisID()); } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/core/tasks/AttachToAnalysisTask.java b/src/main/java/ai/reveng/toolkit/ghidra/core/tasks/AttachToAnalysisTask.java new file mode 100644 index 00000000..3b3da8fe --- /dev/null +++ b/src/main/java/ai/reveng/toolkit/ghidra/core/tasks/AttachToAnalysisTask.java @@ -0,0 +1,68 @@ +package ai.reveng.toolkit.ghidra.core.tasks; + +import ai.reveng.toolkit.ghidra.core.RevEngAIAnalysisResultsLoaded; +import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; +import ghidra.framework.plugintool.PluginTool; +import ghidra.program.model.listing.Function; +import ghidra.util.exception.CancelledException; +import ghidra.util.task.Task; +import ghidra.util.task.TaskMonitor; + +import javax.annotation.Nullable; +import java.util.List; + +/** + * Task that handles attaching to an existing analysis. + * This involves fetching function information from the server and mapping function IDs, + * which can take non-trivial time for large binaries. + */ +public class AttachToAnalysisTask extends Task { + + private final GhidraRevengService.ProgramWithID programWithID; + private final GhidraRevengService service; + private final PluginTool tool; + @Nullable + private final List selectedFunctions; + + /** + * Creates a task to attach to an existing analysis. + * + * @param programWithID The program with associated analysis ID + * @param selectedFunctions Optional list of functions to include in mapping. If null, all functions are included. + * @param service The RevEng.AI service + * @param tool The plugin tool for firing events + */ + public AttachToAnalysisTask( + GhidraRevengService.ProgramWithID programWithID, + @Nullable List selectedFunctions, + GhidraRevengService service, + PluginTool tool + ) { + super("Attaching to RevEng.AI Analysis", false, true, false); + this.programWithID = programWithID; + this.selectedFunctions = selectedFunctions; + this.service = service; + this.tool = tool; + } + + @Override + public void run(TaskMonitor monitor) throws CancelledException { + monitor.setMessage("Fetching function information from server..."); + monitor.setIndeterminate(false); + + var analysedProgram = service.registerFinishedAnalysisForProgram( + programWithID, + selectedFunctions, + monitor + ); + + monitor.setMessage("Analysis attached successfully"); + + tool.firePluginEvent( + new RevEngAIAnalysisResultsLoaded( + "AttachToAnalysisTask", + analysedProgram + ) + ); + } +} \ No newline at end of file diff --git a/src/main/java/ai/reveng/toolkit/ghidra/plugins/AnalysisManagementPlugin.java b/src/main/java/ai/reveng/toolkit/ghidra/plugins/AnalysisManagementPlugin.java index 2573bf71..e6f98158 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/plugins/AnalysisManagementPlugin.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/plugins/AnalysisManagementPlugin.java @@ -26,12 +26,18 @@ import ai.reveng.toolkit.ghidra.core.services.function.export.ExportFunctionBoundariesService; import ai.reveng.toolkit.ghidra.core.services.function.export.ExportFunctionBoundariesServiceImpl; import ai.reveng.toolkit.ghidra.core.services.logging.ReaiLoggingService; +import ai.reveng.toolkit.ghidra.core.tasks.AttachToAnalysisTask; import ai.reveng.toolkit.ghidra.core.tasks.StartAnalysisTask; +import docking.ActionContext; +import docking.ComponentProvider; import docking.action.DockingAction; +import docking.action.MenuData; +import docking.action.ToolBarData; import docking.action.builder.ActionBuilder; import docking.widgets.OptionDialog; import ghidra.app.plugin.PluginCategoryNames; import ghidra.app.plugin.ProgramPlugin; +import ghidra.app.plugin.core.decompile.DecompilerActionContext; import ghidra.app.services.ProgramManager; import ghidra.framework.plugintool.*; import docking.options.OptionsService; @@ -44,6 +50,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import javax.swing.*; import java.util.Objects; /** @@ -113,10 +120,78 @@ public void init() { revengService = Objects.requireNonNull(tool.getService(GhidraRevengService.class)); setupActions(); + installDecompilerToolbarAction(); loggingService.info("CorePlugin initialized"); } + /** + * Installs a toolbar action into the Decompiler window's local toolbar. + * The action is enabled only when the current function is associated with + * an analysis (i.e. has a known FunctionID on the server). + */ + private void installDecompilerToolbarAction() { + ComponentProvider decompilerProvider = tool.getComponentProvider("Decompiler"); + if (decompilerProvider == null) { + Msg.warn(this, "Decompiler provider not found, skipping toolbar action installation"); + return; + } + + DockingAction action = new DockingAction("RevEng.AI Actions", getName()) { + @Override + public boolean isValidContext(ActionContext context) { + return context instanceof DecompilerActionContext; + } + + @Override + public boolean isEnabledForContext(ActionContext context) { + if (!(context instanceof DecompilerActionContext ctx)) { + setDescription("Open function in portal"); + return false; + } + if (ctx.isDecompiling() || !ctx.hasRealFunction()) { + setDescription("Open function in portal"); + return false; + } + var program = ctx.getProgram(); + if (program == null) { + setDescription("Open function in portal"); + return false; + } + var analysedProgram = revengService.getAnalysedProgram(program); + if (analysedProgram.isEmpty()) { + setDescription("No analysis associated with this program"); + return false; + } + var function = ctx.getFunction(); + if (analysedProgram.get().getIDForFunction(function).isEmpty()) { + setDescription("Function has no associated remote function ID"); + return false; + } + setDescription("Open function in portal"); + return true; + } + + @Override + public void actionPerformed(ActionContext context) { + if (!(context instanceof DecompilerActionContext ctx)) { + return; + } + var program = ctx.getProgram(); + if (program == null) { + return; + } + var analysedProgram = revengService.getAnalysedProgram(program).get(); + var function = ctx.getFunction(); + revengService.openPortalFor(analysedProgram.getIDForFunction(function).get()); + } + }; + action.setToolBarData(new ToolBarData(ReaiPluginPackage.REVENG_16, "ZZ_RevEngAI")); + action.setDescription("Open function in portal"); + + tool.addLocalAction(decompilerProvider, action); + } + private void setupActions() { @@ -140,10 +215,10 @@ private void setupActions() { return; } var ghidraService = tool.getService(GhidraRevengService.class); - var dialog = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, ghidraService); + var dialog = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, ghidraService, tool); tool.showDialog(dialog); var analysisOptions = dialog.getOptionsFromUI(); - if (analysisOptions != null) { + if (dialog.isOkPressed()) { // User clicked OK // Prepare Task that starts the analysis (uploading the binary and registering the analysis) var task = new StartAnalysisTask(program, analysisOptions, revengService, analysisLogComponent, tool); @@ -247,23 +322,37 @@ private void setupActions() { .menuGroup(REAI_ANALYSIS_MANAGEMENT_MENU_GROUP, "400") .buildAndInstall(tool); - viewInPortalAction = new ActionBuilder("View in portal", this.getName()) - .enabledWhen(context -> { - var currentProgram = tool.getService(ProgramManager.class).getCurrentProgram(); - if (currentProgram == null) { - // Disable the action if no program is open - return false; - } - return revengService.getKnownProgram(currentProgram).isPresent(); - }) - .onAction(context -> { - var currentProgram = tool.getService(ProgramManager.class).getCurrentProgram(); - var knownProgram = revengService.getKnownProgram(currentProgram).orElseThrow(); - revengService.openPortalFor(knownProgram); - }) - .menuPath(new String[] { ReaiPluginPackage.MENU_GROUP_NAME, "Analysis", "View in portal" }) - .menuGroup(REAI_ANALYSIS_MANAGEMENT_MENU_GROUP, "400") - .buildAndInstall(tool); + viewInPortalAction = new DockingAction("View in portal", getName()) { + @Override + public boolean isEnabledForContext(ActionContext context) { + var currentProgram = tool.getService(ProgramManager.class).getCurrentProgram(); + if (currentProgram == null) { + setDescription("No program is open"); + return false; + } + if (revengService.getKnownProgram(currentProgram).isEmpty()) { + setDescription("No analysis associated with this program"); + return false; + } + setDescription("View analysis in RevEng.AI portal"); + return true; + } + + @Override + public void actionPerformed(ActionContext context) { + var currentProgram = tool.getService(ProgramManager.class).getCurrentProgram(); + var knownProgram = revengService.getKnownProgram(currentProgram).orElseThrow(); + revengService.openPortalFor(knownProgram); + } + }; + viewInPortalAction.setToolBarData(new ToolBarData(ReaiPluginPackage.REVENG_16, "ZZ_RevEngAI")); + viewInPortalAction.setMenuBarData( + new MenuData( + new String[] { ReaiPluginPackage.MENU_GROUP_NAME, "Analysis", "View in portal" }, + null, + REAI_ANALYSIS_MANAGEMENT_MENU_GROUP)); + viewInPortalAction.setDescription("View analysis in RevEng.AI portal"); + tool.addAction(viewInPortalAction); } @Override @@ -333,15 +422,12 @@ public void processEvent(PluginEvent event) { // If the analysis is complete, we refresh the function signatures from the server var program = analysisEvent.getProgramWithBinaryID(); - try { - // TODO: Can we get a better taskmonitor here? - // Or should we never do something here that warrants a monitor in the first place? - var analysedProgram = revengService.registerFinishedAnalysisForProgram(program, TaskMonitor.DUMMY); - tool.firePluginEvent(new RevEngAIAnalysisResultsLoaded("AnalysisManagementPlugin", analysedProgram)); - } catch (Exception e) { - Msg.error(this, "Error registering finished analysis for program " + program, e); - return; - } + var task = new AttachToAnalysisTask(program, null, revengService, tool); + TaskBuilder.withTask(task) + .setCanCancel(false) + .setStatusTextAlignment(SwingConstants.LEADING) + .launchModal(); + } } } diff --git a/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java b/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java index d5b8fe88..665cb022 100644 --- a/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java +++ b/src/main/java/ai/reveng/toolkit/ghidra/plugins/BinarySimilarityPlugin.java @@ -183,9 +183,7 @@ private void setupActions() { .enabledWhen(context -> { var func = context.getProgram().getFunctionManager().getFunctionContaining(context.getAddress()); return func != null - // Exclude thunks and external functions because we do not support them in the portal - && !func.isExternal() - && !func.isThunk() + && GhidraRevengService.isRelevantForAnalysis(func) && apiService.getAnalysedProgram(context.getProgram()).isPresent(); }) .onAction(context -> { @@ -208,9 +206,7 @@ private void setupActions() { .enabledWhen(context -> { var func = context.getProgram().getFunctionManager().getFunctionContaining(context.getAddress()); return func != null - // Exclude thunks and external functions because we do not support them in the portal - && !func.isExternal() - && !func.isThunk() + && GhidraRevengService.isRelevantForAnalysis(func) && apiService.getAnalysedProgram(context.getProgram()).isPresent(); }) .onAction(context -> { diff --git a/src/test/java/ai/reveng/AnalysisOptionsDialogTest.java b/src/test/java/ai/reveng/AnalysisOptionsDialogTest.java index 6d9e0cb3..93cb19e6 100644 --- a/src/test/java/ai/reveng/AnalysisOptionsDialogTest.java +++ b/src/test/java/ai/reveng/AnalysisOptionsDialogTest.java @@ -17,23 +17,24 @@ import static org.junit.Assert.*; +import java.awt.*; +import java.util.*; +import java.util.List; + import javax.swing.*; import ai.reveng.toolkit.ghidra.binarysimilarity.ui.analysiscreation.RevEngAIAnalysisOptionsDialog; +import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection.FunctionSelectionPanel; import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; import ai.reveng.toolkit.ghidra.core.services.api.mocks.MockApi; import docking.DockingWindowManager; -import ghidra.framework.main.FrontEndTool; +import ghidra.program.model.listing.Function; import org.junit.*; import ghidra.program.database.ProgramBuilder; -import ghidra.test.TestEnv; public class AnalysisOptionsDialogTest extends RevEngMockableHeadedIntegrationTest { - private TestEnv env; - private FrontEndTool frontEndTool; - public AnalysisOptionsDialogTest() { super(); } @@ -43,8 +44,21 @@ public void testBasicOptionsDialog() throws Exception { var reService = new GhidraRevengService( new MockApi() {}); var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + // Add some functions to the program so the function selection panel has data + builder.createMemory(".text", "0x401000", 0x1000); + builder.createFunction("0x401000"); + builder.createFunction("0x401100"); + var program = builder.getProgram(); - var dialog = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, reService); + var tool = env.getTool(); + + // Create dialog in the EDT since it uses ThreadedTableModel + var dialogHolder = new RevEngAIAnalysisOptionsDialog[1]; + runSwing(() -> { + dialogHolder[0] = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, reService, tool); + }); + var dialog = dialogHolder[0]; + SwingUtilities.invokeLater(() -> { DockingWindowManager.showDialog(null, dialog); }); @@ -63,4 +77,213 @@ public void testBasicOptionsDialog() throws Exception { capture(dialog.getComponent(), "upload-dialog"); assertNotNull(options); } + + @Test + public void testDialogHasFunctionSelectionPanel() throws Exception { + var reService = new GhidraRevengService(new MockApi() {}); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x401000", 0x1000); + builder.createFunction("0x401000"); + builder.createFunction("0x401100"); + builder.createFunction("0x401200"); + + var program = builder.getProgram(); + var tool = env.getTool(); + + var dialogHolder = new RevEngAIAnalysisOptionsDialog[1]; + runSwing(() -> { + dialogHolder[0] = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, reService, tool); + }); + var dialog = dialogHolder[0]; + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + waitForSwing(); + + // Verify the function selection panel exists + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", dialog); + assertNotNull("Function selection panel should exist", functionSelectionPanel); + + close(dialog); + } + + @Test + public void testFunctionSelectionDefaultsToNonExternalNonThunk() throws Exception { + var reService = new GhidraRevengService(new MockApi() {}); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x401000", 0x1000); + builder.createFunction("0x401000"); + builder.createFunction("0x401100"); + // Create external function (extAddress, libName, functionName) + builder.createExternalFunction(null, "EXTERNAL", "printf"); + + var program = builder.getProgram(); + var tool = env.getTool(); + + var dialogHolder = new RevEngAIAnalysisOptionsDialog[1]; + runSwing(() -> { + dialogHolder[0] = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, reService, tool); + }); + var dialog = dialogHolder[0]; + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + waitForSwing(); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", dialog); + + // Wait for table to load + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // Get selected functions - should not include external + List selectedFunctions = functionSelectionPanel.getSelectedFunctions(); + for (Function func : selectedFunctions) { + assertFalse("External functions should not be selected by default: " + func.getName(), + func.isExternal()); + } + + close(dialog); + } + + @Test + public void testSelectAllButtonWorksInDialog() throws Exception { + var reService = new GhidraRevengService(new MockApi() {}); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x401000", 0x1000); + builder.createFunction("0x401000"); + builder.createFunction("0x401100"); + builder.createFunction("0x401200"); + + var program = builder.getProgram(); + var tool = env.getTool(); + + var dialogHolder = new RevEngAIAnalysisOptionsDialog[1]; + runSwing(() -> { + dialogHolder[0] = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, reService, tool); + }); + var dialog = dialogHolder[0]; + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + waitForSwing(); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", dialog); + + // Wait for table to load + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // Find and click "Select All" button + JButton selectAllButton = findButtonByText(dialog.getComponent(), "Select All"); + assertNotNull("Select All button should exist", selectAllButton); + + pressButton(selectAllButton); + waitForSwing(); + + // All functions should now be selected + assertEquals("All functions should be selected", + functionSelectionPanel.getTotalFunctionCount(), + functionSelectionPanel.getSelectedFunctions().size()); + + close(dialog); + } + + @Test + public void testDeselectAllButtonWorksInDialog() throws Exception { + var reService = new GhidraRevengService(new MockApi() {}); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x401000", 0x1000); + builder.createFunction("0x401000"); + builder.createFunction("0x401100"); + + var program = builder.getProgram(); + var tool = env.getTool(); + + var dialogHolder = new RevEngAIAnalysisOptionsDialog[1]; + runSwing(() -> { + dialogHolder[0] = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, reService, tool); + }); + var dialog = dialogHolder[0]; + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + waitForSwing(); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", dialog); + + // Wait for table to load + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // Find and click "Deselect All" button + JButton deselectAllButton = findButtonByText(dialog.getComponent(), "Deselect All"); + assertNotNull("Deselect All button should exist", deselectAllButton); + + pressButton(deselectAllButton); + waitForSwing(); + + // No functions should be selected + assertTrue("No functions should be selected", + functionSelectionPanel.getSelectedFunctions().isEmpty()); + + close(dialog); + } + + @Test + public void testGetOptionsFromUIIncludesSelectedFunctions() throws Exception { + var reService = new GhidraRevengService(new MockApi() {}); + var builder = new ProgramBuilder("mock", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x401000", 0x1000); + builder.createFunction("0x401000"); + builder.createFunction("0x401100"); + builder.createFunction("0x401200"); + + var program = builder.getProgram(); + var tool = env.getTool(); + + var dialogHolder = new RevEngAIAnalysisOptionsDialog[1]; + runSwing(() -> { + dialogHolder[0] = RevEngAIAnalysisOptionsDialog.withModelsFromServer(program, reService, tool); + }); + var dialog = dialogHolder[0]; + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + waitForSwing(); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", dialog); + + // Wait for table to load + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // By default, all non-external functions should be selected + int selectedCount = functionSelectionPanel.getSelectedFunctions().size(); + assertTrue("Should have at least 3 functions selected", selectedCount >= 3); + + // Get options - should include function boundaries from selected functions + var options = dialog.getOptionsFromUI(); + assertNotNull("Options should be returned", options); + + // Convert to AnalysisCreateRequest to inspect the function boundaries + var request = options.toAnalysisCreateRequest(); + assertNotNull("Request should be created", request); + assertNotNull("Request should have symbols", request.getSymbols()); + assertNotNull("Symbols should have function boundaries", request.getSymbols().getFunctionBoundaries()); + + var boundaries = request.getSymbols().getFunctionBoundaries(); + assertEquals("Function boundaries count should match selected functions", + selectedCount, boundaries.size()); + + // Verify the expected function addresses are present (0x401000, 0x401100, 0x401200) + var startAddresses = boundaries.stream() + .map(b -> b.getStartAddress()) + .toList(); + assertTrue("Should contain function at 0x401000", startAddresses.contains(0x401000L)); + assertTrue("Should contain function at 0x401100", startAddresses.contains(0x401100L)); + assertTrue("Should contain function at 0x401200", startAddresses.contains(0x401200L)); + + close(dialog); + } } diff --git a/src/test/java/ai/reveng/RecentAnalysisDialogTest.java b/src/test/java/ai/reveng/RecentAnalysisDialogTest.java index 46715d3b..da8c0823 100644 --- a/src/test/java/ai/reveng/RecentAnalysisDialogTest.java +++ b/src/test/java/ai/reveng/RecentAnalysisDialogTest.java @@ -1,18 +1,26 @@ package ai.reveng; +import ai.reveng.toolkit.ghidra.binarysimilarity.ui.functionselection.FunctionSelectionPanel; import ai.reveng.toolkit.ghidra.binarysimilarity.ui.recentanalyses.RecentAnalysisDialog; +import ai.reveng.toolkit.ghidra.core.RevEngAIAnalysisResultsLoaded; import ai.reveng.toolkit.ghidra.core.RevEngAIAnalysisStatusChangedEvent; import ai.reveng.toolkit.ghidra.core.services.api.GhidraRevengService; import ai.reveng.toolkit.ghidra.core.services.api.TypedApiInterface; import ai.reveng.toolkit.ghidra.core.services.api.mocks.UnimplementedAPI; import ai.reveng.toolkit.ghidra.core.services.api.types.AnalysisStatus; import ai.reveng.toolkit.ghidra.core.services.api.types.BinaryID; +import ai.reveng.toolkit.ghidra.core.services.api.types.FunctionInfo; import ai.reveng.toolkit.ghidra.core.services.api.types.LegacyAnalysisResult; +import ai.reveng.model.FunctionDataTypesList; import docking.DockingWindowManager; import ghidra.program.database.ProgramBuilder; +import ghidra.program.model.data.Undefined; +import ghidra.program.model.data.Undefined4DataType; +import ghidra.program.model.listing.Function; import org.junit.Test; import javax.swing.*; +import java.awt.*; import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; @@ -33,21 +41,22 @@ public void testSelectRecentAnalysisFiresEventAndUpdatesKnownProgram() throws Ex var mockApi = new RecentAnalysesMockApi(); var service = addMockedService(tool, mockApi); - // Create a test program with matching hash + // Create a test program with matching hash and at least one function var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); - builder.createMemory("test", "0x1000", 100); + builder.createMemory(".text", "0x1000", 0x1000); + builder.createEmptyFunction("main", "0x1000", 100, Undefined.getUndefinedDataType(4)); var program = builder.getProgram(); // Show the tool with the program env.showTool(program); waitForSwing(); - // Set up event listener to capture the analysis status changed event + // Set up event listener to capture the analysis results loaded event AtomicBoolean eventReceived = new AtomicBoolean(false); - AtomicReference receivedEvent = new AtomicReference<>(); - tool.addEventListener(RevEngAIAnalysisStatusChangedEvent.class, e -> { + AtomicReference receivedEvent = new AtomicReference<>(); + tool.addEventListener(RevEngAIAnalysisResultsLoaded.class, e -> { eventReceived.set(true); - receivedEvent.set((RevEngAIAnalysisStatusChangedEvent) e); + receivedEvent.set((RevEngAIAnalysisResultsLoaded) e); }); // Verify the program is not known before the dialog interaction @@ -90,23 +99,21 @@ public void testSelectRecentAnalysisFiresEventAndUpdatesKnownProgram() throws Ex // Verify the event was fired with correct data assertNotNull("Event should have been captured", receivedEvent.get()); - assertEquals("Event status should match the analysis status", - AnalysisStatus.Complete, receivedEvent.get().getStatus()); - GhidraRevengService.ProgramWithID eventProgramWithID = receivedEvent.get().getProgramWithBinaryID(); - assertNotNull("Event should contain ProgramWithID", eventProgramWithID); + GhidraRevengService.AnalysedProgram analysedProgram = receivedEvent.get().getProgramWithBinaryID(); + assertNotNull("Event should contain AnalysedProgram", analysedProgram); assertSame("Event program should be the same as our test program", - program, eventProgramWithID.program()); + program, analysedProgram.program()); assertEquals("Event analysis ID should match mock data", - RecentAnalysesMockApi.MOCK_ANALYSIS_ID, eventProgramWithID.analysisID().id()); + RecentAnalysesMockApi.MOCK_ANALYSIS_ID, analysedProgram.analysisID().id()); // Verify getKnownProgram returns the same program with the correct analysis ID var knownProgram = service.getKnownProgram(program); assertTrue("Program should be known after selection", knownProgram.isPresent()); assertEquals("Known program analysis ID should match event analysis ID", - eventProgramWithID.analysisID(), knownProgram.get().analysisID()); + analysedProgram.analysisID(), knownProgram.get().analysisID()); assertSame("Known program should be the same instance as event program", - eventProgramWithID.program(), knownProgram.get().program()); + analysedProgram.program(), knownProgram.get().program()); } @Test @@ -185,5 +192,393 @@ public AnalysisStatus status(TypedApiInterface.AnalysisID analysisID) { assertEquals("Analysis ID should match mock data", MOCK_ANALYSIS_ID, analysisID.id()); return AnalysisStatus.Complete; } + + @Override + public List getFunctionInfo(TypedApiInterface.AnalysisID analysisID) { + // Return function info covering common test addresses. + // Note: 0x1200 is deliberately omitted so tests can verify unmatched function behavior. + return List.of( + new FunctionInfo( + new TypedApiInterface.FunctionID(1001), + "main", + "main", + 0x1000L, + 100 + ), + new FunctionInfo( + new TypedApiInterface.FunctionID(1002), + "helper", + "helper", + 0x1100L, + 50 + ) + ); + } + + @Override + public FunctionDataTypesList listFunctionDataTypesForAnalysis(TypedApiInterface.AnalysisID analysisID) { + // Return empty list + var list = new FunctionDataTypesList(); + list.setItems(List.of()); + return list; + } + } + + // ==================== Function Selection Tests ==================== + + @Test + public void testDialogHasFunctionSelectionPanel() throws Exception { + var tool = env.getTool(); + + var mockApi = new RecentAnalysesMockApi(); + addMockedService(tool, mockApi); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x1000", 0x1000); + builder.createEmptyFunction("main", "0x1000", 100, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction("helper", "0x1100", 50, Undefined.getUndefinedDataType(4)); + var program = builder.getProgram(); + + env.showTool(program); + waitForSwing(); + + RecentAnalysisDialog dialog = runSwing(() -> + new RecentAnalysisDialog(tool, program) + ); + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(RecentAnalysisDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Verify the function selection panel exists + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", foundDialog); + assertNotNull("Function selection panel should exist in RecentAnalysisDialog", + functionSelectionPanel); + + close(foundDialog); + waitForSwing(); + } + + @Test + public void testFunctionSelectionPanelLoadsInAttachDialog() throws Exception { + var tool = env.getTool(); + + var mockApi = new RecentAnalysesMockApi(); + addMockedService(tool, mockApi); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x1000", 0x1000); + builder.createEmptyFunction("main", "0x1000", 100, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction("process", "0x1100", 150, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction("cleanup", "0x1200", 80, Undefined.getUndefinedDataType(4)); + var program = builder.getProgram(); + + env.showTool(program); + waitForSwing(); + + RecentAnalysisDialog dialog = runSwing(() -> + new RecentAnalysisDialog(tool, program) + ); + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(RecentAnalysisDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", foundDialog); + + // Wait for functions to load + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // Should have at least our 3 functions + assertTrue("Should have at least 3 functions", + functionSelectionPanel.getTotalFunctionCount() >= 3); + + close(foundDialog); + waitForSwing(); + } + + @Test + public void testSelectAllButtonInAttachDialog() throws Exception { + var tool = env.getTool(); + + var mockApi = new RecentAnalysesMockApi(); + addMockedService(tool, mockApi); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x1000", 0x1000); + builder.createEmptyFunction("func1", "0x1000", 100, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction("func2", "0x1100", 100, Undefined.getUndefinedDataType(4)); + var program = builder.getProgram(); + + env.showTool(program); + waitForSwing(); + + RecentAnalysisDialog dialog = runSwing(() -> + new RecentAnalysisDialog(tool, program) + ); + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(RecentAnalysisDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", foundDialog); + + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // Find and click Select All button + JButton selectAllButton = findButtonByText(foundDialog.getComponent(), "Select All"); + assertNotNull("Select All button should exist", selectAllButton); + + pressButton(selectAllButton); + waitForSwing(); + + // All enabled functions should be selected (disabled functions from remote matching remain unselected) + int enabledCount = functionSelectionPanel.getTableModel().getEnabledCount(); + assertEquals("All enabled functions should be selected", + enabledCount, + functionSelectionPanel.getSelectedFunctions().size()); + assertTrue("At least one function should be selected", enabledCount > 0); + + close(foundDialog); + waitForSwing(); + } + + @Test + public void testDeselectAllButtonInAttachDialog() throws Exception { + var tool = env.getTool(); + + var mockApi = new RecentAnalysesMockApi(); + addMockedService(tool, mockApi); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x1000", 0x1000); + builder.createEmptyFunction("func1", "0x1000", 100, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction("func2", "0x1100", 100, Undefined.getUndefinedDataType(4)); + var program = builder.getProgram(); + + env.showTool(program); + waitForSwing(); + + RecentAnalysisDialog dialog = runSwing(() -> + new RecentAnalysisDialog(tool, program) + ); + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(RecentAnalysisDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Wait for analyses table to load so the auto-select triggers the remote function fetch + var tableModelField = getInstanceField("recentAnalysesTableModel", foundDialog); + @SuppressWarnings("unchecked") + var tableModel = (docking.widgets.table.threaded.ThreadedTableModel) tableModelField; + waitForTableModel(tableModel); + waitForSwing(); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", foundDialog); + + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // Wait for remote function matching to complete (async fetch) + waitForCondition(() -> functionSelectionPanel.getTableModel().getEnabledCount() < functionSelectionPanel.getTotalFunctionCount(), + "Remote function matching should complete before testing deselect"); + + // Find and click Deselect All button + JButton deselectAllButton = findButtonByText(foundDialog.getComponent(), "Deselect All"); + assertNotNull("Deselect All button should exist", deselectAllButton); + + pressButton(deselectAllButton); + waitForSwing(); + + // No functions should be selected + assertTrue("No functions should be selected", + functionSelectionPanel.getSelectedFunctions().isEmpty()); + + close(foundDialog); + waitForSwing(); + } + + @Test + public void testExternalFunctionsExcludedByDefaultInAttachDialog() throws Exception { + var tool = env.getTool(); + + var mockApi = new RecentAnalysesMockApi(); + addMockedService(tool, mockApi); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x1000", 0x1000); + builder.createEmptyFunction("main", "0x1000", 100, Undefined.getUndefinedDataType(4)); + // Create external function (extAddress, libName, functionName) + builder.createExternalFunction(null, "EXTERNAL", "printf"); + var program = builder.getProgram(); + + env.showTool(program); + waitForSwing(); + + RecentAnalysisDialog dialog = runSwing(() -> + new RecentAnalysisDialog(tool, program) + ); + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(RecentAnalysisDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", foundDialog); + + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // External functions should not be selected by default + List selectedFunctions = functionSelectionPanel.getSelectedFunctions(); + for (Function func : selectedFunctions) { + assertFalse("External function should not be selected by default: " + func.getName(), + func.isExternal()); + } + + close(foundDialog); + waitForSwing(); } + + @Test + public void testRemoteFunctionMatchingDisablesUnmatchedFunctions() throws Exception { + var tool = env.getTool(); + + var mockApi = new RecentAnalysesMockApi(); + addMockedService(tool, mockApi); + + // Create 3 functions: main (0x1000, 100), helper (0x1100, 50), unmatched (0x1200, 80) + // The mock API returns info for main and helper but not unmatched + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x1000", 0x1000); + builder.createEmptyFunction("main", "0x1000", 100, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction("helper", "0x1100", 50, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction("unmatched", "0x1200", 80, Undefined.getUndefinedDataType(4)); + var program = builder.getProgram(); + + env.showTool(program); + waitForSwing(); + + RecentAnalysisDialog dialog = runSwing(() -> + new RecentAnalysisDialog(tool, program) + ); + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(RecentAnalysisDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Wait for the analyses table to load + var tableModelField = getInstanceField("recentAnalysesTableModel", foundDialog); + @SuppressWarnings("unchecked") + var tableModel = (docking.widgets.table.threaded.ThreadedTableModel) tableModelField; + waitForTableModel(tableModel); + waitForSwing(); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", foundDialog); + + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // Wait for remote function info to be applied (async fetch) + waitForCondition(() -> functionSelectionPanel.getTableModel().getEnabledCount() < functionSelectionPanel.getTotalFunctionCount(), + "Remote function matching should disable some functions"); + + // Verify enabled/disabled state + var allRows = functionSelectionPanel.getTableModel().getAllRows(); + for (var row : allRows) { + if ("unmatched".equals(row.getName())) { + assertFalse("Unmatched function should be disabled", row.isEnabled()); + assertFalse("Unmatched function should not be selected", row.isSelected()); + assertNull("Unmatched function should have no remote info", row.getRemoteFunctionInfo()); + } else { + assertTrue("Matched function '" + row.getName() + "' should be enabled", row.isEnabled()); + assertNotNull("Matched function '" + row.getName() + "' should have remote info", row.getRemoteFunctionInfo()); + } + } + + // Verify that select all only selects enabled functions + JButton selectAllButton = findButtonByText(foundDialog.getComponent(), "Select All"); + pressButton(selectAllButton); + waitForSwing(); + + for (var row : allRows) { + if ("unmatched".equals(row.getName())) { + assertFalse("Disabled function should remain unselected after Select All", row.isSelected()); + } else { + assertTrue("Enabled function should be selected after Select All", row.isSelected()); + } + } + + close(foundDialog); + waitForSwing(); + } + + @Test + public void testPickAnalysisWithFunctionSelection() throws Exception { + var tool = env.getTool(); + + var mockApi = new RecentAnalysesMockApi(); + var service = addMockedService(tool, mockApi); + + var builder = new ProgramBuilder("test_binary", ProgramBuilder._X64, this); + builder.createMemory(".text", "0x1000", 0x1000); + builder.createEmptyFunction("main", "0x1000", 100, Undefined.getUndefinedDataType(4)); + builder.createEmptyFunction("helper", "0x1100", 50, Undefined.getUndefinedDataType(4)); + var program = builder.getProgram(); + + env.showTool(program); + waitForSwing(); + + // Set up event listener + AtomicBoolean eventReceived = new AtomicBoolean(false); + tool.addEventListener(RevEngAIAnalysisStatusChangedEvent.class, e -> { + eventReceived.set(true); + }); + + RecentAnalysisDialog dialog = runSwing(() -> + new RecentAnalysisDialog(tool, program) + ); + + runSwing(() -> DockingWindowManager.showDialog(null, dialog), false); + var foundDialog = waitForDialogComponent(RecentAnalysisDialog.class); + assertNotNull("Dialog should be shown", foundDialog); + + // Wait for both table models to load + var tableModelField = getInstanceField("recentAnalysesTableModel", foundDialog); + @SuppressWarnings("unchecked") + var tableModel = (docking.widgets.table.threaded.ThreadedTableModel) tableModelField; + waitForTableModel(tableModel); + + FunctionSelectionPanel functionSelectionPanel = + (FunctionSelectionPanel) getInstanceField("functionSelectionPanel", foundDialog); + waitForCondition(() -> functionSelectionPanel.getTotalFunctionCount() > 0, + "Function selection panel should load functions"); + + // Select only the main function (deselect helper) + JButton deselectAllButton = findButtonByText(foundDialog.getComponent(), "Deselect All"); + pressButton(deselectAllButton); + waitForSwing(); + + // With no functions selected, the pick buttons should be disabled + JButton pickMostRecentButton = findButtonByText(foundDialog.getComponent(), "Pick most recent"); + assertNotNull("Pick most recent button should exist", pickMostRecentButton); + assertFalse("Pick most recent button should be disabled when no functions are selected", + pickMostRecentButton.isEnabled()); + + JButton pickSelectedButton = findButtonByText(foundDialog.getComponent(), "Pick selected"); + assertNotNull("Pick selected button should exist", pickSelectedButton); + assertFalse("Pick selected button should be disabled when no functions are selected", + pickSelectedButton.isEnabled()); + + close(foundDialog); + waitForSwing(); + } + } \ No newline at end of file diff --git a/src/test/java/ai/reveng/RevEngMockableHeadedIntegrationTest.java b/src/test/java/ai/reveng/RevEngMockableHeadedIntegrationTest.java index 650d9ae4..f7bd2cc9 100644 --- a/src/test/java/ai/reveng/RevEngMockableHeadedIntegrationTest.java +++ b/src/test/java/ai/reveng/RevEngMockableHeadedIntegrationTest.java @@ -14,11 +14,11 @@ import java.io.IOException; /// Base class -abstract class RevEngMockableHeadedIntegrationTest extends AbstractGhidraHeadedIntegrationTest { +abstract public class RevEngMockableHeadedIntegrationTest extends AbstractGhidraHeadedIntegrationTest { protected TestEnv env; @Before - public void setup() throws IOException, PluginException { + public void setup() throws Exception { // For most tests we want to fail if a user visible error would show up. // Ghidra already has a nifty feature for that, we just need to activate it setErrorGUIEnabled(false);