Add entry functions in triage view

3 files changed, 362 insertions, 1 deletions

diff --git a/examples/triage/entry.cpp b/examples/triage/entry.cpp
new file mode 100644
index 00000000..4c7f5ec0
--- /dev/null
+++ b/examples/triage/entry.cpp
@@ -0,0 +1,281 @@
+#include <cstring>
+#include <algorithm>
+#include "entry.h"
+#include "view.h"
+#include "fontsettings.h"
+
+
+const int AddressColumn = 0;
+const int NameColumn = 1;
+const int ColumnCount = 2;
+
+
+GenericEntryModel::GenericEntryModel(QWidget* parent, BinaryViewRef data): QAbstractItemModel(parent)
+{
+	m_data = data;
+	m_sortOrder = Qt::AscendingOrder;
+	m_allEntries = data->GetAllEntryFunctions();
+	m_entries = m_allEntries;
+}
+
+
+int GenericEntryModel::columnCount(const QModelIndex&) const
+{
+	return ColumnCount;
+}
+
+
+int GenericEntryModel::rowCount(const QModelIndex& parent) const
+{
+	if (parent.isValid())
+		return 0;
+	return (int)m_entries.size();
+}
+
+
+QVariant GenericEntryModel::data(const QModelIndex& index, int role) const
+{
+	switch (role)
+	{
+	case Qt::DisplayRole:
+		if (!index.isValid() || index.row() >= (int)m_entries.size())
+			return QVariant();
+		if (index.column() == AddressColumn)
+			return QString("0x") + QString::number(m_entries[index.row()]->GetStart(), 16);
+		if (index.column() == NameColumn)
+			return QString::fromStdString(m_entries[index.row()]->GetSymbol()->GetFullName());
+		break;
+	case Qt::ForegroundRole:
+		if (index.column() == AddressColumn)
+			return getThemeColor(AddressColor);
+		if (index.column() == NameColumn)
+			return getThemeColor(ExportColor);
+		break;
+	default:
+		break;
+	}
+
+	return QVariant();
+}
+
+
+QVariant GenericEntryModel::headerData(int section, Qt::Orientation orientation, int role) const
+{
+	if (orientation == Qt::Vertical)
+		return QVariant();
+	if (role != Qt::DisplayRole)
+		return QVariant();
+	if (section == AddressColumn)
+		return QString("Address");
+	if (section == NameColumn)
+		return QString("Name");
+	return QVariant();
+}
+
+
+QModelIndex GenericEntryModel::index(int row, int col, const QModelIndex& parent) const
+{
+	if (parent.isValid())
+		return QModelIndex();
+	if (row >= (int)m_entries.size())
+		return QModelIndex();
+	if (col >= ColumnCount)
+		return QModelIndex();
+	return createIndex(row, col);
+}
+
+
+QModelIndex GenericEntryModel::parent(const QModelIndex&) const
+{
+	return QModelIndex();
+}
+
+
+FunctionRef GenericEntryModel::getEntry(const QModelIndex& index)
+{
+	if (!index.isValid() || index.row() >= (int)m_entries.size())
+		return nullptr;
+	return m_entries[index.row()];
+}
+
+
+void GenericEntryModel::performSort(int col, Qt::SortOrder order)
+{
+	std::sort(m_entries.begin(), m_entries.end(), [&](FunctionRef a, FunctionRef b) {
+		if (col == AddressColumn)
+		{
+			if (a->GetStart() != b->GetStart())
+			{
+				if (order == Qt::AscendingOrder)
+					return a->GetStart() < b->GetStart();
+				else
+					return a->GetStart() > b->GetStart();
+			}
+			if (order == Qt::AscendingOrder)
+				return a->GetSymbol()->GetFullName() < b->GetSymbol()->GetFullName();
+			else
+				return a->GetSymbol()->GetFullName() > b->GetSymbol()->GetFullName();
+		}
+		else if (col == NameColumn)
+		{
+			if (order == Qt::AscendingOrder)
+				return a->GetSymbol()->GetFullName() < b->GetSymbol()->GetFullName();
+			else
+				return a->GetSymbol()->GetFullName() > b->GetSymbol()->GetFullName();
+		}
+		return false;
+	});
+}
+
+
+void GenericEntryModel::sort(int col, Qt::SortOrder order)
+{
+	beginResetModel();
+	m_sortCol = col;
+	m_sortOrder = order;
+	performSort(col, order);
+	endResetModel();
+}
+
+
+void GenericEntryModel::setFilter(const std::string& filterText)
+{
+	beginResetModel();
+	m_entries.clear();
+	for (auto& entry : m_allEntries)
+	{
+		if (FilteredView::match(entry->GetSymbol()->GetFullName(), filterText))
+			m_entries.push_back(entry);
+	}
+	performSort(m_sortCol, m_sortOrder);
+	endResetModel();
+}
+
+
+EntryTreeView::EntryTreeView(EntryWidget* parent, TriageView* view, BinaryViewRef data) : QTreeView(parent)
+{
+	m_data = data;
+	m_parent = parent;
+	m_view = view;
+
+	// Allow view-specific shortcuts when imports are focused
+	m_actionHandler.setupActionHandler(this);
+	m_actionHandler.setActionContext([=]() { return m_view->actionContext(); });
+
+	m_model = new GenericEntryModel(this, m_data);
+	setModel(m_model);
+	setRootIsDecorated(false);
+	setUniformRowHeights(true);
+	setSortingEnabled(true);
+	sortByColumn(AddressColumn, Qt::AscendingOrder);
+
+	setColumnWidth(AddressColumn, 90);
+
+	setFont(getMonospaceFont(this));
+
+	connect(selectionModel(), &QItemSelectionModel::currentChanged, this, &EntryTreeView::entrySelected);
+	connect(this, &QTreeView::doubleClicked, this, &EntryTreeView::entryDoubleClicked);
+}
+
+
+void EntryTreeView::entrySelected(const QModelIndex& cur, const QModelIndex&)
+{
+	FunctionRef func = m_model->getEntry(cur);
+	if (func)
+		m_view->setCurrentOffset(func->GetStart());
+}
+
+
+void EntryTreeView::entryDoubleClicked(const QModelIndex& cur)
+{
+	FunctionRef func = m_model->getEntry(cur);
+	if (func)
+	{
+		ViewFrame* viewFrame = ViewFrame::viewFrameForWidget(this);
+		if (viewFrame)
+		{
+			if (BinaryNinja::Settings::Instance()->Get<bool>("ui.view.graph.preferred") &&
+				viewFrame->getCurrentBinaryView() &&
+				func->GetStart() > 0)
+			{
+				viewFrame->navigate("Graph:" + viewFrame->getCurrentDataType(), func->GetStart());
+			}
+			else
+			{
+				viewFrame->navigate("Linear:" + viewFrame->getCurrentDataType(), func->GetStart());
+			}
+		}
+	}
+}
+
+
+void EntryTreeView::setFilter(const std::string& filterText)
+{
+	m_model->setFilter(filterText);
+}
+
+
+void EntryTreeView::scrollToFirstItem()
+{
+	scrollToTop();
+}
+
+
+void EntryTreeView::scrollToCurrentItem()
+{
+	scrollTo(currentIndex());
+}
+
+
+void EntryTreeView::selectFirstItem()
+{
+	setCurrentIndex(m_model->index(0, 0, QModelIndex()));
+}
+
+
+void EntryTreeView::activateFirstItem()
+{
+	entryDoubleClicked(m_model->index(0, 0, QModelIndex()));
+}
+
+
+void EntryTreeView::closeFilter()
+{
+	setFocus(Qt::OtherFocusReason);
+}
+
+
+void EntryTreeView::keyPressEvent(QKeyEvent* event)
+{
+	if ((event->text().size() == 1) && (event->text()[0] > ' ') && (event->text()[0] <= '~'))
+	{
+		m_parent->showFilter(event->text());
+		event->accept();
+	}
+	else if ((event->key() == Qt::Key_Return) || (event->key() == Qt::Key_Enter))
+	{
+		QList<QModelIndex> sel = selectionModel()->selectedIndexes();
+		if (sel.size() != 0)
+			entryDoubleClicked(sel[0]);
+	}
+	QTreeView::keyPressEvent(event);
+}
+
+
+EntryWidget::EntryWidget(QWidget* parent, TriageView* view, BinaryViewRef data) : QWidget(parent)
+{
+	QVBoxLayout* layout = new QVBoxLayout();
+	layout->setContentsMargins(0, 0, 0, 0);
+	EntryTreeView* entry = new EntryTreeView(this, view, data);
+	m_filter = new FilteredView(this, entry, entry);
+	m_filter->setFilterPlaceholderText("Search entry functions");
+	layout->addWidget(m_filter, 1);
+	setLayout(layout);
+	setMinimumSize(UIContext::getScaledWindowSize(100, 196));
+}
+
+
+void EntryWidget::showFilter(const QString& filter)
+{
+	m_filter->showFilter(filter);
+}
diff --git a/examples/triage/entry.h b/examples/triage/entry.h
new file mode 100644
index 00000000..3ed82db5
--- /dev/null
+++ b/examples/triage/entry.h
@@ -0,0 +1,70 @@
+#pragma once
+
+#include <QtCore/QAbstractItemModel>
+#include <QtWidgets/QTreeView>
+#include "filter.h"
+
+
+class GenericEntryModel : public QAbstractItemModel
+{
+	BinaryViewRef m_data;
+	std::vector<FunctionRef> m_allEntries, m_entries;
+	Qt::SortOrder m_sortOrder;
+	int m_sortCol;
+
+	void performSort(int col, Qt::SortOrder order);
+
+  public:
+	GenericEntryModel(QWidget* parent, BinaryViewRef data);
+
+	virtual int columnCount(const QModelIndex& parent) const override;
+	virtual int rowCount(const QModelIndex& parent) const override;
+	virtual QVariant data(const QModelIndex& index, int role) const override;
+	virtual QVariant headerData(int section, Qt::Orientation orientation, int role) const override;
+	virtual QModelIndex index(int row, int col, const QModelIndex& parent) const override;
+	virtual QModelIndex parent(const QModelIndex& index) const override;
+	virtual void sort(int col, Qt::SortOrder order) override;
+	void setFilter(const std::string& filterText);
+
+	FunctionRef getEntry(const QModelIndex& index);
+};
+
+
+class TriageView;
+class EntryWidget;
+
+class EntryTreeView : public QTreeView, public FilterTarget
+{
+	BinaryViewRef m_data;
+	EntryWidget* m_parent;
+	TriageView* m_view;
+	UIActionHandler m_actionHandler;
+	GenericEntryModel* m_model;
+
+  public:
+	EntryTreeView(EntryWidget* parent, TriageView* view, BinaryViewRef data);
+
+	virtual void setFilter(const std::string& filterText) override;
+	virtual void scrollToFirstItem() override;
+	virtual void scrollToCurrentItem() override;
+	virtual void selectFirstItem() override;
+	virtual void activateFirstItem() override;
+	virtual void closeFilter() override;
+
+  protected:
+	virtual void keyPressEvent(QKeyEvent* event) override;
+
+  private Q_SLOTS:
+	void entrySelected(const QModelIndex& cur, const QModelIndex& prev);
+	void entryDoubleClicked(const QModelIndex& cur);
+};
+
+
+class EntryWidget : public QWidget
+{
+	FilteredView* m_filter;
+
+  public:
+	EntryWidget(QWidget* parent, TriageView* view, BinaryViewRef data);
+	void showFilter(const QString& filter);
+};
diff --git a/examples/triage/view.cpp b/examples/triage/view.cpp
index ad44da8c..ee12e8a0 100644
--- a/examples/triage/view.cpp
+++ b/examples/triage/view.cpp
@@ -3,6 +3,7 @@
 #include <QtWidgets/QSplitter>
 #include "view.h"
 #include "entropy.h"
+#include "entry.h"
 #include "imports.h"
 #include "exports.h"
 #include "sections.h"
@@ -80,12 +81,21 @@ TriageView::TriageView(QWidget* parent, BinaryViewRef data) : QScrollArea(parent
 		importGroup->setLayout(importLayout);
 		importExportSplitter->addWidget(importGroup);
 
+		QSplitter* exportEntrySplitter = new QSplitter(Qt::Vertical);
+
 		QGroupBox* exportGroup = new QGroupBox("Exports", container);
 		QVBoxLayout* exportLayout = new QVBoxLayout();
 		exportLayout->addWidget(new ExportsWidget(exportGroup, this, m_data));
 		exportGroup->setLayout(exportLayout);
-		importExportSplitter->addWidget(exportGroup);
+		exportEntrySplitter->addWidget(exportGroup);
+
+		QGroupBox* entryGroup = new QGroupBox("Entry Functions", container);
+		QVBoxLayout* entryLayout = new QVBoxLayout();
+		entryLayout->addWidget(new EntryWidget(entryGroup, this, m_data));
+		entryGroup->setLayout(entryLayout);
+		exportEntrySplitter->addWidget(entryGroup);
 
+		importExportSplitter->addWidget(exportEntrySplitter);
 		layout->addWidget(importExportSplitter);
 
 		if (m_data->GetTypeName() != "PE")