kaitai UI plugin + elf,macho,pe

1 files changed, 220 insertions, 0 deletions

diff --git a/python/examples/kaitai/view.py b/python/examples/kaitai/view.py
new file mode 100644
index 00000000..888049ac
--- /dev/null
+++ b/python/examples/kaitai/view.py
@@ -0,0 +1,220 @@
+import traceback
+import binaryninjaui
+from binaryninja.settings import Settings
+from binaryninja import log
+from binaryninja import _binaryninjacore as core
+from binaryninjaui import View, ViewType, ViewFrame, UIContext, HexEditor
+from binaryninja import binaryview
+from PySide2.QtWidgets import QScrollArea, QWidget, QVBoxLayout, QGroupBox, QTreeWidget, QTreeWidgetItem, QLineEdit
+from PySide2.QtCore import Qt
+
+import types
+
+import kshelpers
+from kaitaistruct import KaitaiStruct
+import io
+
+class KaitaiView(QScrollArea, View):
+	def __init__(self, parent, binaryView):
+		QScrollArea.__init__(self, parent)
+
+		View.__init__(self)
+		self.setupView(self)
+
+		# BinaryViewType
+		self.binaryView = binaryView
+
+		self.rootSelectionStart = 0
+		self.rootSelectionEnd = 1
+
+		self.ioRoot = None
+		self.ioCurrent = None
+
+		container = QWidget(self)
+		layout = QVBoxLayout()
+
+		self.treeGroup = QGroupBox("Data Tree:")
+		treeLayout = QVBoxLayout()
+		self.treeWidget = QTreeWidget()
+		self.treeWidget.setColumnCount(4)
+		self.treeWidget.setHeaderLabels(['label','value','start','end'])
+		self.treeWidget.itemSelectionChanged.connect(self.onTreeSelect)
+		self.structPath = QLineEdit("root")
+		self.structPath.setDisabled(True)
+		treeLayout.addWidget(self.structPath)
+		treeLayout.addWidget(self.treeWidget)
+		self.treeGroup.setLayout(treeLayout)
+
+		self.hexGroup = QGroupBox("Hex View:")
+		self.hexLayout = QVBoxLayout()
+		self.hexWidget = HexEditor(binaryView, ViewFrame.viewFrameForWidget(self), 0)
+		self.hexLayout.addWidget(self.hexWidget)
+		self.hexGroup.setLayout(self.hexLayout)
+
+		layout.addWidget(self.treeGroup)
+		layout.addWidget(self.hexGroup)
+		#layout.addStretch(1)
+		container.setLayout(layout)
+		self.setWidgetResizable(True)
+		self.setWidget(container)
+
+		self.kaitaiParse()
+
+	# parse the file using Kaitai, construct the TreeWidget
+	def kaitaiParse(self):
+		parsed = None
+
+		kaitaiIO = kshelpers.KaitaiBinaryViewIO(self.binaryView)
+		if not kaitaiIO:
+			print 'ERROR: initializing kaitai binary view'
+		parsed = kshelpers.parse_io(kaitaiIO)
+		if not parsed:
+			print 'ERROR: parsing the binary view'
+
+		tree = kshelpers.build_qtree(parsed)
+		if not tree:
+			return
+
+		self.ioRoot = tree.ksobj._io
+		self.ioCurrent = tree.ksobj._io
+
+		self.treeWidget.clear()
+		self.treeWidget.setSortingEnabled(False)						# temporarily, for efficiency
+		# two options with how we create the hierarchy
+		if False:
+			# treat root as top level "file" container
+			tree.setLabel('file')
+			tree.setValue(None)
+			tree.setStart(0)
+			tree.setEnd(0)
+			self.treeWidget.insertTopLevelItem(0, tree)
+		else:
+			# add root's children as top level items
+			self.treeWidget.insertTopLevelItems(0, tree.takeChildren())
+
+		#self.treeWidget.expandAll()
+		width = self.treeWidget.width()
+		self.treeWidget.setColumnWidth(0, .4*width)
+		self.treeWidget.setColumnWidth(1, .2*width)
+		self.treeWidget.setColumnWidth(2, .1*width)
+		self.treeWidget.setColumnWidth(3, .1*width)
+		# enable sorting
+		self.treeWidget.setSortingEnabled(True)
+		self.treeWidget.sortByColumn(2, Qt.AscendingOrder)
+
+		# TODO: select first item, maybe expand a few things
+		self.rootSelectionStart = 0
+		self.rootSelectionEnd = 1
+		self.hexWidget.setSelectionRange(0,1)
+		self.treeWidget.setUniformRowHeights(True)
+
+	# callback!
+	def getData(self):
+		return self.binaryView
+
+	# callback!
+	def getCurrentOffset(self):
+		middle = self.rootSelectionStart + int((self.rootSelectionEnd - self.rootSelectionStart)/2)
+		return middle
+
+	def setCurrentOffset(self, offset):
+		print 'setCurrentOffset(0x%X)' % offset
+		self.rootSelectionStart = offset
+		UIContext.updateStatus(True)
+
+	def getFont(self):
+		return binaryninjaui.getMonospaceFont(self)
+
+	def navigate(self, addr):
+		return False
+
+	def onTreeSelect(self, wtf=None):
+		# get KaitaiTreeWidgetItem
+		item = self.treeWidget.selectedItems()[0]
+
+		# build path, inform user
+		structPath = item.label
+		itemTmp = item
+		while itemTmp.parent():
+			itemTmp = itemTmp.parent()
+			structPath = itemTmp.label + '.' + structPath
+		self.structPath.setText('root.' + structPath)
+
+		#
+		(start, end) = (item.start, item.end)
+		if start == None or end == None:
+			return
+
+		# determine current IO we're in (the Kaitai input/output abstraction)
+		_io = None
+		# if the tree item is linked to a KaitaiNode, simply read the IO
+		if item.ksobj:
+			_io = item.ksobj._parent._io
+		else:
+			# else we're a leaf
+			parent = item.parent()
+			if parent:
+				# a leaf with a parent -> read parent's IO
+				_io = parent.ksobj._io
+			else:
+				# a leaf without a parent -> we must be at root -> use root IO
+				_io = self.ioRoot
+
+		# if the selection is in the root view, store the interval so that upon
+		# getCurrentOffset() callback, we return the middle and feature map is
+		# updated
+		if _io == self.ioRoot:
+			self.rootSelectionStart = start
+			self.rootSelectionEnd = end
+
+		# current kaitai object is on a different io? then swap HexEditor
+		if _io != self.ioCurrent:
+			# delete old view
+			layoutItem = self.hexLayout.takeAt(0)
+			hexEditorWidget = layoutItem.widget()
+			hexEditorWidget.setParent(None)
+			hexEditorWidget.deleteLater()
+			self.hexWidget = None
+
+			# if it's the original file IO, wrap the already-open file binary view
+			if _io == self.ioRoot:
+				self.hexWidget = HexEditor(self.binaryView, ViewFrame.viewFrameForWidget(self), 0)
+			# otherwise delete old view, create a temporary view
+			else:
+				# create new view
+				length = _io.size()
+				_io.seek(0)
+				data = _io.read_bytes(length)
+				bv = binaryview.BinaryView.new(data)
+				self.hexWidget = HexEditor(bv, ViewFrame.viewFrameForWidget(self), 0)
+
+			self.hexLayout.addWidget(self.hexWidget)
+			self.ioCurrent = _io
+
+		# now position selection in whatever HexEditor is current
+		#print 'selecting to [0x%X, 0x%X)' % (start, end)
+		self.hexWidget.setSelectionRange(start, end)
+
+		# set hex group title to reflect current selection
+		self.hexGroup.setTitle('Hex View @ [0x%X, 0x%X)' % (start, end))
+
+class KaitaiViewType(ViewType):
+	def __init__(self):
+		super(KaitaiViewType, self).__init__("Kaitai", "Kaitai")
+
+	# binaryView:		BinaryView
+	def getPriority(self, binaryView, filename):
+		# data.file:	FileMetadata
+		# data.raw:		BinaryView
+
+		if binaryView == binaryView.file.raw:
+			#print 'returning priority=100 for FileView'
+			return 100
+		else:
+			#print 'returning priority 0, not given a raw view'
+			return 0
+
+	def create(self, binaryView, view_frame):
+		return KaitaiView(view_frame, binaryView)
+
+ViewType.registerViewType(KaitaiViewType())