Add a tokenized text view class for UI plugins

3 files changed, 308 insertions, 0 deletions

diff --git a/binaryninjacore.h b/binaryninjacore.h
index c4aac243..9eb74713 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -260,6 +260,7 @@ extern "C"
 	enum BNLinearDisassemblyLineType
 	{
 		BlankLineType,
+		BasicLineType,
 		CodeDisassemblyLineType,
 		DataVariableLineType,
 		HexDumpLineType,
diff --git a/python/examples/linear_mlil.py b/python/examples/linear_mlil.py
new file mode 100644
index 00000000..403dddd2
--- /dev/null
+++ b/python/examples/linear_mlil.py
@@ -0,0 +1,139 @@
+# Copyright (c) 2019 Vector 35 Inc
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+
+from binaryninja.function import DisassemblyTextRenderer, DisassemblyTextLine
+from binaryninja.lineardisassembly import LinearDisassemblyLine
+from binaryninja.enums import LinearDisassemblyLineType, DisassemblyOption
+from binaryninjaui import TokenizedTextView, TokenizedTextViewHistoryEntry, ViewType
+
+
+class LinearMLILView(TokenizedTextView):
+	def __init__(self, parent, data):
+		super(LinearMLILView, self).__init__(parent, data)
+		self.data = data
+		self.function = data.entry_function
+		if self.function is not None:
+			self.setFunction(self.function)
+			self.updateLines()
+
+	def generateLines(self):
+		if self.function is None:
+			return []
+
+		il = self.function.mlil
+
+		# Set up IL display options
+		renderer = DisassemblyTextRenderer(il)
+		renderer.settings.set_option(DisassemblyOption.ShowAddress)
+		renderer.settings.set_option(DisassemblyOption.ShowVariableTypesWhenAssigned)
+
+		# Sort basic blocks by IL instruction index
+		blocks = il.basic_blocks
+		blocks.sort(key = lambda block: block.start)
+
+		# Function header
+		result = []
+		result.append(LinearDisassemblyLine(LinearDisassemblyLineType.FunctionHeaderStartLineType,
+			self.function, None, 0, DisassemblyTextLine([], self.function.start)))
+		result.append(LinearDisassemblyLine(LinearDisassemblyLineType.FunctionHeaderLineType,
+			self.function, None, 0, DisassemblyTextLine(self.function.type_tokens, self.function.start)))
+		result.append(LinearDisassemblyLine(LinearDisassemblyLineType.FunctionHeaderEndLineType,
+			self.function, None, 0, DisassemblyTextLine([], self.function.start)))
+
+		# Display IL instructions in order
+		lastAddr = self.function.start
+		lastBlock = None
+		lineIndex = 0
+		for block in il:
+			if lastBlock is not None:
+				# Blank line between basic blocks
+				result.append(LinearDisassemblyLine(LinearDisassemblyLineType.CodeDisassemblyLineType,
+					self.function, block, 0, DisassemblyTextLine([], lastAddr)))
+			for i in block:
+				lines, length = renderer.get_disassembly_text(i.instr_index)
+				lastAddr = i.address
+				lineIndex = 0
+				for line in lines:
+					result.append(LinearDisassemblyLine(LinearDisassemblyLineType.CodeDisassemblyLineType,
+						self.function, block, lineIndex, line))
+					lineIndex += 1
+			lastBlock = block
+
+		result.append(LinearDisassemblyLine(LinearDisassemblyLineType.FunctionEndLineType,
+			self.function, lastBlock, lineIndex, DisassemblyTextLine([], lastAddr)))
+
+		return result
+
+	def updateLines(self):
+		self.setUpdatedLines(self.generateLines())
+
+	def navigate(self, addr):
+		# Find correct function based on most recent use
+		block = self.data.get_recent_basic_block_at(addr)
+		if block is None:
+			# If function isn't done analyzing yet, it may have a function start but no basic blocks
+			func = self.data.get_recent_function_at(addr)
+		else:
+			func = block.function
+
+		if func is None:
+			# No function contains this address, fail navigation in this view
+			return False
+
+		self.function = func
+		self.setFunction(self.function)
+		self.setLines(self.generateLines())
+		return True
+
+	def getHistoryEntry(self):
+		class LinearMLILHistoryEntry(TokenizedTextViewHistoryEntry):
+			def __init__(self, function):
+				super(LinearMLILHistoryEntry, self).__init__()
+				self.function = function
+
+		result = LinearMLILHistoryEntry(self.function)
+		self.populateDefaultHistoryEntry(result)
+		return result
+
+	def navigateToHistoryEntry(self, entry):
+		if hasattr(entry, 'function'):
+			self.function = entry.function
+			self.setFunction(self.function)
+			self.updateLines()
+		super(LinearMLILView, self).navigateToHistoryEntry(entry)
+
+
+# View type for the new view
+class LinearMLILViewType(ViewType):
+	def __init__(self):
+		super(LinearMLILViewType, self).__init__("Linear MLIL", "Linear MLIL")
+
+	def getPriority(self, data, filename):
+		if data.executable:
+			# Use low priority so that this view is not picked by default
+			return 1
+		return 0
+
+	def create(self, data, view_frame):
+		return LinearMLILView(view_frame, data)
+
+
+# Register the view type so that it can be chosen by the user
+ViewType.registerViewType(LinearMLILViewType())
diff --git a/ui/tokenizedtextview.h b/ui/tokenizedtextview.h
new file mode 100644
index 00000000..2a64d971
--- /dev/null
+++ b/ui/tokenizedtextview.h
@@ -0,0 +1,168 @@
+#pragma once
+
+#include <QtWidgets/QAbstractScrollArea>
+#include <QtCore/QTimer>
+#include "binaryninjaapi.h"
+#include "viewframe.h"
+#include "render.h"
+#include "commentdialog.h"
+#include "menus.h"
+#include "uicontext.h"
+
+class BINARYNINJAUIAPI TokenizedTextViewHistoryEntry: public HistoryEntry
+{
+	size_t m_topLine, m_cursorLine;
+	HighlightTokenState m_highlight;
+
+public:
+	size_t getTopLine() const { return m_topLine; }
+	size_t getCursorLine() const { return m_cursorLine; }
+	const HighlightTokenState& getHighlightTokenState() const { return m_highlight; }
+
+	void setTopLine(size_t line) { m_topLine = line; }
+	void setCursorLine(size_t line) { m_cursorLine = line; }
+	void setHighlightTokenState(const HighlightTokenState& state) { m_highlight = state; }
+};
+
+class BINARYNINJAUIAPI TokenizedTextView: public QAbstractScrollArea, public View, public BinaryNinja::BinaryDataNotification
+{
+	Q_OBJECT
+
+	BinaryViewRef m_data;
+	FunctionRef m_function;
+
+	RenderContext m_render;
+	int m_cols, m_rows;
+	int m_wheelDelta;
+	bool m_updatingScrollBar;
+
+	bool m_updatesRequired;
+
+	int m_cursorLine;
+	HighlightTokenState m_highlight;
+	uint64_t m_navByRefTarget;
+	bool m_navByRef = false;
+
+	std::vector<BinaryNinja::LinearDisassemblyLine> m_lines;
+
+	QTimer* m_updateTimer;
+
+	ContextMenuManager m_contextMenuManager;
+	QPointer<CommentDialog> m_commentDialog;
+
+	void adjustSize(int width, int height);
+
+	void scrollLines(int count);
+
+	void bindActions();
+	void getHexDumpLineBytes(const BinaryNinja::LinearDisassemblyLine& line, size_t& skippedBytes, size_t& totalBytes,
+		size_t& totalCols);
+
+	void setSectionSemantics(const std::string& name, BNSectionSemantics semantics);
+
+	void viewInHexEditor();
+	void viewInGraph();
+	void viewInTypesView(std::string typeName = "");
+	void goToAddress();
+	void defineNameAtAddr(uint64_t addr);
+	void defineName();
+	void undefineName();
+	void createFunc();
+	void defineFuncName();
+	void undefineFunc();
+	void reanalyze();
+	void comment();
+	void commentAccepted();
+
+	void convertToNop();
+	void alwaysBranch();
+	void invertBranch();
+	void skipAndReturnZero();
+	void skipAndReturnValue();
+
+	void makeInt8();
+	void makeInt16();
+	void makeInt32();
+	void makeInt64();
+	void toggleIntSize();
+	void makePtr();
+	void makeString();
+	void changeType();
+	size_t getStringLength(uint64_t startAddr);
+
+	void displayAsDefault();
+	void displayAsBinary();
+	void displayAsSignedOctal();
+	void displayAsUnsignedOctal();
+	void displayAsSignedDecimal();
+	void displayAsUnsignedDecimal();
+	void displayAsSignedHexadecimal();
+	void displayAsUnsignedHexadecimal();
+	void displayAsCharacterConstant();
+	void displayAsPointer();
+
+	void setInstructionHighlight(BNHighlightColor color);
+	void setBlockHighlight(BNHighlightColor color);
+
+private Q_SLOTS:
+	void scrollBarMoved(int value);
+	void scrollBarAction(int action);
+	void updateTimerEvent();
+
+public:
+	explicit TokenizedTextView(QWidget* parent, BinaryViewRef data,
+		const std::vector<BinaryNinja::LinearDisassemblyLine>& lines = std::vector<BinaryNinja::LinearDisassemblyLine>());
+	virtual ~TokenizedTextView();
+
+	virtual BinaryViewRef getData() override { return m_data; }
+	virtual uint64_t getCurrentOffset() override;
+	virtual void getSelectionOffsets(uint64_t& begin, uint64_t& end) override;
+	virtual void getSelectionForInfo(uint64_t& begin, uint64_t& end) override;
+	virtual FunctionRef getCurrentFunction() override;
+	virtual BasicBlockRef getCurrentBasicBlock() override;
+	virtual ArchitectureRef getCurrentArchitecture() override;
+	virtual bool navigate(uint64_t pos) override;
+
+	virtual HistoryEntry* getHistoryEntry() override;
+	void populateDefaultHistoryEntry(TokenizedTextViewHistoryEntry* entry);
+	virtual void navigateToHistoryEntry(HistoryEntry* entry) override;
+
+	virtual void OnBinaryDataWritten(BinaryNinja::BinaryView* data, uint64_t offset, size_t len) override;
+	virtual void OnBinaryDataInserted(BinaryNinja::BinaryView* data, uint64_t offset, size_t len) override;
+	virtual void OnBinaryDataRemoved(BinaryNinja::BinaryView* data, uint64_t offset, uint64_t len) override;
+	virtual void OnAnalysisFunctionAdded(BinaryNinja::BinaryView* view, BinaryNinja::Function* func) override;
+	virtual void OnAnalysisFunctionRemoved(BinaryNinja::BinaryView* view, BinaryNinja::Function* func) override;
+	virtual void OnAnalysisFunctionUpdated(BinaryNinja::BinaryView* view, BinaryNinja::Function* func) override;
+	virtual void OnDataVariableAdded(BinaryNinja::BinaryView* view, const BinaryNinja::DataVariable& var) override;
+	virtual void OnDataVariableRemoved(BinaryNinja::BinaryView* view, const BinaryNinja::DataVariable& var) override;
+	virtual void OnDataVariableUpdated(BinaryNinja::BinaryView* view, const BinaryNinja::DataVariable& var) override;
+
+	virtual void updateFonts() override;
+
+	virtual void followPointer();
+
+	virtual void cut() override;
+	virtual void copy(TransformRef xform = nullptr) override;
+	virtual void paste(TransformRef xform = nullptr) override;
+	virtual void copyAddress() override;
+
+	virtual HighlightTokenState getHighlightTokenState() override { return m_highlight; }
+
+	virtual bool goToReference(FunctionRef func, uint64_t source, uint64_t target) override;
+	QFont getFont() override { return m_render.getFont(); }
+
+	static void registerActions();
+
+	virtual void updateLines();
+	void setLines(const std::vector<BinaryNinja::LinearDisassemblyLine>& lines);
+	void setUpdatedLines(const std::vector<BinaryNinja::LinearDisassemblyLine>& lines);
+
+	void setFunction(FunctionRef func);
+
+protected:
+	virtual void resizeEvent(QResizeEvent* event) override;
+	virtual void paintEvent(QPaintEvent* event) override;
+	virtual void wheelEvent(QWheelEvent* event) override;
+	virtual void mousePressEvent(QMouseEvent* event) override;
+	virtual void mouseDoubleClickEvent(QMouseEvent* event) override;
+};