mirror of
https://github.com/mandiant/capa.git
synced 2025-12-22 07:10:29 -08:00
Merge pull request #58 from fireeye/capa-explorer-support-doc-format
Capa explorer support doc format
This commit is contained in:
@@ -1,4 +1,3 @@
|
||||
import binascii
|
||||
import codecs
|
||||
import sys
|
||||
|
||||
@@ -10,115 +9,116 @@ import idc
|
||||
import capa.ida.helpers
|
||||
|
||||
|
||||
def info_to_name(s):
|
||||
''' '''
|
||||
def info_to_name(display):
|
||||
""" extract root value from display name
|
||||
|
||||
e.g. function(my_function) => my_function
|
||||
"""
|
||||
try:
|
||||
return s.split('(')[1].rstrip(')')
|
||||
return display.split('(')[1].rstrip(')')
|
||||
except IndexError:
|
||||
return ''
|
||||
|
||||
|
||||
def ea_to_hex_str(ea):
|
||||
''' '''
|
||||
return '%08X' % ea
|
||||
def location_to_hex(location):
|
||||
""" convert location to hex for display """
|
||||
return '%08X' % location
|
||||
|
||||
|
||||
class CapaExplorerDataItem(object):
|
||||
''' store data for CapaExplorerDataModel
|
||||
""" store data for CapaExplorerDataModel """
|
||||
|
||||
TODO
|
||||
'''
|
||||
def __init__(self, parent, data):
|
||||
''' '''
|
||||
self._parent = parent
|
||||
""" """
|
||||
self.pred = parent
|
||||
self._data = data
|
||||
self._children = []
|
||||
self.children = []
|
||||
self._checked = False
|
||||
|
||||
self.flags = (QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsTristate | QtCore.Qt.ItemIsUserCheckable)
|
||||
|
||||
if self._parent:
|
||||
self._parent.appendChild(self)
|
||||
if self.pred:
|
||||
self.pred.appendChild(self)
|
||||
|
||||
def setIsEditable(self, isEditable=False):
|
||||
''' modify item flags to be editable or not '''
|
||||
""" modify item flags to be editable or not """
|
||||
if isEditable:
|
||||
self.flags |= QtCore.Qt.ItemIsEditable
|
||||
else:
|
||||
self.flags &= ~QtCore.Qt.ItemIsEditable
|
||||
|
||||
def setChecked(self, checked):
|
||||
''' set item as checked '''
|
||||
""" set item as checked """
|
||||
self._checked = checked
|
||||
|
||||
def isChecked(self):
|
||||
''' get item is checked '''
|
||||
""" get item is checked """
|
||||
return self._checked
|
||||
|
||||
def appendChild(self, item):
|
||||
''' add child item
|
||||
""" add child item
|
||||
|
||||
@param item: CapaExplorerDataItem*
|
||||
'''
|
||||
self._children.append(item)
|
||||
"""
|
||||
self.children.append(item)
|
||||
|
||||
def child(self, row):
|
||||
''' get child row
|
||||
""" get child row
|
||||
|
||||
@param row: TODO
|
||||
'''
|
||||
return self._children[row]
|
||||
"""
|
||||
return self.children[row]
|
||||
|
||||
def childCount(self):
|
||||
''' get child count '''
|
||||
return len(self._children)
|
||||
""" get child count """
|
||||
return len(self.children)
|
||||
|
||||
def columnCount(self):
|
||||
''' get column count '''
|
||||
""" get column count """
|
||||
return len(self._data)
|
||||
|
||||
def data(self, column):
|
||||
''' get data at column '''
|
||||
""" get data at column """
|
||||
try:
|
||||
return self._data[column]
|
||||
except IndexError:
|
||||
return None
|
||||
|
||||
def parent(self):
|
||||
''' get parent '''
|
||||
return self._parent
|
||||
""" get parent """
|
||||
return self.pred
|
||||
|
||||
def row(self):
|
||||
''' get row location '''
|
||||
if self._parent:
|
||||
return self._parent._children.index(self)
|
||||
""" get row location """
|
||||
if self.pred:
|
||||
return self.pred.children.index(self)
|
||||
return 0
|
||||
|
||||
def setData(self, column, value):
|
||||
''' set data in column '''
|
||||
""" set data in column """
|
||||
self._data[column] = value
|
||||
|
||||
def children(self):
|
||||
''' yield children '''
|
||||
for child in self._children:
|
||||
""" yield children """
|
||||
for child in self.children:
|
||||
yield child
|
||||
|
||||
def removeChildren(self):
|
||||
''' '''
|
||||
del self._children[:]
|
||||
""" remove children from node """
|
||||
del self.children[:]
|
||||
|
||||
def __str__(self):
|
||||
''' get string representation of columns '''
|
||||
""" get string representation of columns """
|
||||
return ' '.join([data for data in self._data if data])
|
||||
|
||||
@property
|
||||
def info(self):
|
||||
''' '''
|
||||
""" return data stored in information column """
|
||||
return self._data[0]
|
||||
|
||||
@property
|
||||
def ea(self):
|
||||
''' '''
|
||||
def location(self):
|
||||
""" return data stored in location column """
|
||||
try:
|
||||
return int(self._data[1], 16)
|
||||
except ValueError:
|
||||
@@ -126,107 +126,108 @@ class CapaExplorerDataItem(object):
|
||||
|
||||
@property
|
||||
def details(self):
|
||||
''' '''
|
||||
""" return data stored in details column """
|
||||
return self._data[2]
|
||||
|
||||
|
||||
class CapaExplorerRuleItem(CapaExplorerDataItem):
|
||||
''' store data relevant to capa function result '''
|
||||
""" store data relevant to capa function result """
|
||||
|
||||
view_fmt = '%s (%d)'
|
||||
fmt = '%s (%d matches)'
|
||||
|
||||
def __init__(self, parent, name, count, definition):
|
||||
''' '''
|
||||
self._definition = definition
|
||||
name = CapaExplorerRuleItem.view_fmt % (name, count) if count else name
|
||||
super(CapaExplorerRuleItem, self).__init__(parent, [name, '', ''])
|
||||
def __init__(self, parent, display, count, source):
|
||||
""" """
|
||||
display = self.fmt % (display, count) if count > 1 else display
|
||||
super(CapaExplorerRuleItem, self).__init__(parent, [display, '', ''])
|
||||
self._source = source
|
||||
|
||||
@property
|
||||
def definition(self):
|
||||
''' '''
|
||||
return self._definition
|
||||
def source(self):
|
||||
""" return rule contents for display """
|
||||
return self._source
|
||||
|
||||
|
||||
class CapaExplorerRuleMatchItem(CapaExplorerDataItem):
|
||||
""" store data relevant to capa function match result """
|
||||
|
||||
def __init__(self, parent, display, source=''):
|
||||
""" """
|
||||
super(CapaExplorerRuleMatchItem, self).__init__(parent, [display, '', ''])
|
||||
self._source = source
|
||||
|
||||
@property
|
||||
def source(self):
|
||||
""" return rule contents for display """
|
||||
return self._source
|
||||
|
||||
|
||||
class CapaExplorerFunctionItem(CapaExplorerDataItem):
|
||||
''' store data relevant to capa function result '''
|
||||
""" store data relevant to capa function result """
|
||||
|
||||
view_fmt = 'function(%s)'
|
||||
fmt = 'function(%s)'
|
||||
|
||||
def __init__(self, parent, name, ea):
|
||||
''' '''
|
||||
address = ea_to_hex_str(ea)
|
||||
name = CapaExplorerFunctionItem.view_fmt % name
|
||||
|
||||
super(CapaExplorerFunctionItem, self).__init__(parent, [name, address, ''])
|
||||
def __init__(self, parent, location):
|
||||
""" """
|
||||
super(CapaExplorerFunctionItem, self).__init__(parent, [self.fmt % idaapi.get_name(location),
|
||||
location_to_hex(location), ''])
|
||||
|
||||
@property
|
||||
def info(self):
|
||||
''' '''
|
||||
""" """
|
||||
info = super(CapaExplorerFunctionItem, self).info
|
||||
name = info_to_name(info)
|
||||
return name if name else info
|
||||
display = info_to_name(info)
|
||||
return display if display else info
|
||||
|
||||
@info.setter
|
||||
def info(self, name):
|
||||
''' '''
|
||||
self._data[0] = CapaExplorerFunctionItem.view_fmt % name
|
||||
def info(self, display):
|
||||
""" """
|
||||
self._data[0] = self.fmt % display
|
||||
|
||||
|
||||
class CapaExplorerBlockItem(CapaExplorerDataItem):
|
||||
''' store data relevant to capa basic block results '''
|
||||
""" store data relevant to capa basic block result """
|
||||
|
||||
view_fmt = 'basic block(loc_%s)'
|
||||
fmt = 'basic block(loc_%08X)'
|
||||
|
||||
def __init__(self, parent, ea):
|
||||
''' '''
|
||||
address = ea_to_hex_str(ea)
|
||||
name = CapaExplorerBlockItem.view_fmt % address
|
||||
|
||||
super(CapaExplorerBlockItem, self).__init__(parent, [name, address, ''])
|
||||
def __init__(self, parent, location):
|
||||
""" """
|
||||
super(CapaExplorerBlockItem, self).__init__(parent, [self.fmt % location, location_to_hex(location), ''])
|
||||
|
||||
|
||||
class CapaExplorerDefaultItem(CapaExplorerDataItem):
|
||||
''' store data relevant to capa default result '''
|
||||
""" store data relevant to capa default result """
|
||||
|
||||
def __init__(self, parent, name, ea=None):
|
||||
''' '''
|
||||
if ea:
|
||||
address = ea_to_hex_str(ea)
|
||||
else:
|
||||
address = ''
|
||||
|
||||
super(CapaExplorerDefaultItem, self).__init__(parent, [name, address, ''])
|
||||
def __init__(self, parent, display, details='', location=None):
|
||||
""" """
|
||||
location = location_to_hex(location) if location else ''
|
||||
super(CapaExplorerDefaultItem, self).__init__(parent, [display, location, details])
|
||||
|
||||
|
||||
class CapaExplorerFeatureItem(CapaExplorerDataItem):
|
||||
''' store data relevant to capa feature result '''
|
||||
""" store data relevant to capa feature result """
|
||||
|
||||
def __init__(self, parent, data):
|
||||
super(CapaExplorerFeatureItem, self).__init__(parent, data)
|
||||
def __init__(self, parent, display, location='', details=''):
|
||||
location = location_to_hex(location) if location else ''
|
||||
super(CapaExplorerFeatureItem, self).__init__(parent, [display, location, details])
|
||||
|
||||
|
||||
class CapaExplorerInstructionViewItem(CapaExplorerFeatureItem):
|
||||
|
||||
def __init__(self, parent, name, ea):
|
||||
''' '''
|
||||
details = capa.ida.helpers.get_disasm_line(ea)
|
||||
address = ea_to_hex_str(ea)
|
||||
|
||||
super(CapaExplorerInstructionViewItem, self).__init__(parent, [name, address, details])
|
||||
|
||||
self.ida_highlight = idc.get_color(ea, idc.CIC_ITEM)
|
||||
def __init__(self, parent, display, location):
|
||||
""" """
|
||||
details = capa.ida.helpers.get_disasm_line(location)
|
||||
super(CapaExplorerInstructionViewItem, self).__init__(parent, display, location=location, details=details)
|
||||
self.ida_highlight = idc.get_color(location, idc.CIC_ITEM)
|
||||
|
||||
|
||||
class CapaExplorerByteViewItem(CapaExplorerFeatureItem):
|
||||
|
||||
def __init__(self, parent, name, ea):
|
||||
''' '''
|
||||
address = ea_to_hex_str(ea)
|
||||
def __init__(self, parent, display, location):
|
||||
""" """
|
||||
byte_snap = idaapi.get_bytes(location, 32)
|
||||
|
||||
byte_snap = idaapi.get_bytes(ea, 32)
|
||||
if byte_snap:
|
||||
byte_snap = codecs.encode(byte_snap, 'hex').upper()
|
||||
# TODO: better way?
|
||||
if sys.version_info >= (3, 0):
|
||||
details = ' '.join([byte_snap[i:i + 2].decode() for i in range(0, len(byte_snap), 2)])
|
||||
else:
|
||||
@@ -234,17 +235,13 @@ class CapaExplorerByteViewItem(CapaExplorerFeatureItem):
|
||||
else:
|
||||
details = ''
|
||||
|
||||
super(CapaExplorerByteViewItem, self).__init__(parent, [name, address, details])
|
||||
|
||||
self.ida_highlight = idc.get_color(ea, idc.CIC_ITEM)
|
||||
super(CapaExplorerByteViewItem, self).__init__(parent, display, location=location, details=details)
|
||||
self.ida_highlight = idc.get_color(location, idc.CIC_ITEM)
|
||||
|
||||
|
||||
class CapaExplorerStringViewItem(CapaExplorerFeatureItem):
|
||||
|
||||
def __init__(self, parent, name, ea, value):
|
||||
''' '''
|
||||
address = ea_to_hex_str(ea)
|
||||
|
||||
super(CapaExplorerStringViewItem, self).__init__(parent, [name, address, value])
|
||||
|
||||
self.ida_highlight = idc.get_color(ea, idc.CIC_ITEM)
|
||||
def __init__(self, parent, display, location):
|
||||
""" """
|
||||
super(CapaExplorerStringViewItem, self).__init__(parent, display, location=location)
|
||||
self.ida_highlight = idc.get_color(location, idc.CIC_ITEM)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from PyQt5 import QtCore, QtGui
|
||||
from PyQt5 import QtCore, QtGui, Qt
|
||||
from collections import deque
|
||||
import binascii
|
||||
|
||||
import capa.render.utils as rutils
|
||||
|
||||
import idaapi
|
||||
import idc
|
||||
@@ -8,24 +9,24 @@ import idc
|
||||
from capa.ida.explorer.item import (
|
||||
CapaExplorerDataItem,
|
||||
CapaExplorerDefaultItem,
|
||||
CapaExplorerFeatureItem,
|
||||
CapaExplorerFunctionItem,
|
||||
CapaExplorerRuleItem,
|
||||
CapaExplorerStringViewItem,
|
||||
CapaExplorerInstructionViewItem,
|
||||
CapaExplorerByteViewItem,
|
||||
CapaExplorerBlockItem
|
||||
CapaExplorerBlockItem,
|
||||
CapaExplorerRuleMatchItem,
|
||||
CapaExplorerFeatureItem
|
||||
)
|
||||
|
||||
import capa.ida.helpers
|
||||
|
||||
|
||||
# default highlight color used in IDA window
|
||||
DEFAULT_HIGHLIGHT = 0xD096FF
|
||||
|
||||
|
||||
class CapaExplorerDataModel(QtCore.QAbstractItemModel):
|
||||
''' '''
|
||||
""" """
|
||||
|
||||
COLUMN_INDEX_RULE_INFORMATION = 0
|
||||
COLUMN_INDEX_VIRTUAL_ADDRESS = 1
|
||||
@@ -34,114 +35,134 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
|
||||
COLUMN_COUNT = 3
|
||||
|
||||
def __init__(self, parent=None):
|
||||
''' '''
|
||||
""" """
|
||||
super(CapaExplorerDataModel, self).__init__(parent)
|
||||
|
||||
self._root = CapaExplorerDataItem(None, ['Rule Information', 'Address', 'Details'])
|
||||
self.root_node = CapaExplorerDataItem(None, ['Rule Information', 'Address', 'Details'])
|
||||
|
||||
def reset(self):
|
||||
''' '''
|
||||
""" """
|
||||
# reset checkboxes and color highlights
|
||||
# TODO: make less hacky
|
||||
for idx in range(self._root.childCount()):
|
||||
rindex = self.index(idx, 0, QtCore.QModelIndex())
|
||||
for mindex in self.iterateChildrenIndexFromRootIndex(rindex, ignore_root=False):
|
||||
mindex.internalPointer().setChecked(False)
|
||||
self._util_reset_ida_highlighting(mindex.internalPointer(), False)
|
||||
self.dataChanged.emit(mindex, mindex)
|
||||
for idx in range(self.root_node.childCount()):
|
||||
root_index = self.index(idx, 0, QtCore.QModelIndex())
|
||||
for model_index in self.iterateChildrenIndexFromRootIndex(root_index, ignore_root=False):
|
||||
model_index.internalPointer().setChecked(False)
|
||||
self.util_reset_ida_highlighting(model_index.internalPointer(), False)
|
||||
self.dataChanged.emit(model_index, model_index)
|
||||
|
||||
def clear(self):
|
||||
''' '''
|
||||
""" """
|
||||
self.beginResetModel()
|
||||
# TODO: make sure this isn't for memory
|
||||
self._root.removeChildren()
|
||||
self.root_node.removeChildren()
|
||||
self.endResetModel()
|
||||
|
||||
def columnCount(self, mindex):
|
||||
''' get the number of columns for the children of the given parent
|
||||
def columnCount(self, model_index):
|
||||
""" get the number of columns for the children of the given parent
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
@param model_index: QModelIndex*
|
||||
|
||||
@retval column count
|
||||
'''
|
||||
if mindex.isValid():
|
||||
return mindex.internalPointer().columnCount()
|
||||
"""
|
||||
if model_index.isValid():
|
||||
return model_index.internalPointer().columnCount()
|
||||
else:
|
||||
return self._root.columnCount()
|
||||
return self.root_node.columnCount()
|
||||
|
||||
def data(self, mindex, role):
|
||||
''' get data stored under the given role for the item referred to by the index
|
||||
def data(self, model_index, role):
|
||||
""" get data stored under the given role for the item referred to by the index
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
@param model_index: QModelIndex*
|
||||
@param role: QtCore.Qt.*
|
||||
|
||||
@retval data to be displayed
|
||||
'''
|
||||
if not mindex.isValid():
|
||||
"""
|
||||
if not model_index.isValid():
|
||||
return None
|
||||
|
||||
item = model_index.internalPointer()
|
||||
column = model_index.column()
|
||||
|
||||
if role == QtCore.Qt.DisplayRole:
|
||||
# display data in corresponding column
|
||||
return mindex.internalPointer().data(mindex.column())
|
||||
return item.data(column)
|
||||
|
||||
if role == QtCore.Qt.ToolTipRole and \
|
||||
CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == mindex.column() and \
|
||||
isinstance(mindex.internalPointer(), CapaExplorerRuleItem):
|
||||
# show tooltip containing rule definition
|
||||
return mindex.internalPointer().definition
|
||||
if role == QtCore.Qt.ToolTipRole and isinstance(item, (CapaExplorerRuleItem, CapaExplorerRuleMatchItem)) and \
|
||||
CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == column:
|
||||
# show tooltip containing rule source
|
||||
return item.source
|
||||
|
||||
if role == QtCore.Qt.CheckStateRole and mindex.column() == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION:
|
||||
if role == QtCore.Qt.CheckStateRole and column == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION:
|
||||
# inform view how to display content of checkbox - un/checked
|
||||
return QtCore.Qt.Checked if mindex.internalPointer().isChecked() else QtCore.Qt.Unchecked
|
||||
return QtCore.Qt.Checked if item.isChecked() else QtCore.Qt.Unchecked
|
||||
|
||||
if role == QtCore.Qt.FontRole and mindex.column() in (CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, CapaExplorerDataModel.COLUMN_INDEX_DETAILS):
|
||||
return QtGui.QFont('Courier', weight=QtGui.QFont.Medium)
|
||||
if role == QtCore.Qt.FontRole and column in (CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS,
|
||||
CapaExplorerDataModel.COLUMN_INDEX_DETAILS):
|
||||
# set font for virtual address and details columns
|
||||
font = QtGui.QFont('Courier', weight=QtGui.QFont.Medium)
|
||||
if column == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS:
|
||||
font.setBold(True)
|
||||
return font
|
||||
|
||||
if role == QtCore.Qt.FontRole and mindex.internalPointer() == self._root:
|
||||
return QtCore.QFont(bold=True)
|
||||
if role == QtCore.Qt.FontRole and isinstance(item, (CapaExplorerRuleItem, CapaExplorerRuleMatchItem,
|
||||
CapaExplorerBlockItem, CapaExplorerFunctionItem,
|
||||
CapaExplorerFeatureItem)) and \
|
||||
column == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION:
|
||||
# set bold font for top-level rules
|
||||
font = QtGui.QFont()
|
||||
font.setBold(True)
|
||||
return font
|
||||
|
||||
if role == QtCore.Qt.ForegroundRole and column == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS:
|
||||
# set color for virtual address column
|
||||
return QtGui.QColor(88, 139, 174)
|
||||
|
||||
if role == QtCore.Qt.ForegroundRole and isinstance(item, CapaExplorerFeatureItem) and column == \
|
||||
CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION:
|
||||
# set color for feature items
|
||||
return QtGui.QColor(79, 121, 66)
|
||||
|
||||
return None
|
||||
|
||||
def flags(self, mindex):
|
||||
''' get item flags for given index
|
||||
def flags(self, model_index):
|
||||
""" get item flags for given index
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
@param model_index: QModelIndex*
|
||||
|
||||
@retval QtCore.Qt.ItemFlags
|
||||
'''
|
||||
if not mindex.isValid():
|
||||
"""
|
||||
if not model_index.isValid():
|
||||
return QtCore.Qt.NoItemFlags
|
||||
|
||||
return mindex.internalPointer().flags
|
||||
return model_index.internalPointer().flags
|
||||
|
||||
def headerData(self, section, orientation, role):
|
||||
''' get data for the given role and section in the header with the specified orientation
|
||||
""" get data for the given role and section in the header with the specified orientation
|
||||
|
||||
@param section: int
|
||||
@param orientation: QtCore.Qt.Orientation
|
||||
@param role: QtCore.Qt.DisplayRole
|
||||
|
||||
@retval header data list()
|
||||
'''
|
||||
"""
|
||||
if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
|
||||
return self._root.data(section)
|
||||
return self.root_node.data(section)
|
||||
|
||||
return None
|
||||
|
||||
def index(self, row, column, parent):
|
||||
''' get index of the item in the model specified by the given row, column and parent index
|
||||
""" get index of the item in the model specified by the given row, column and parent index
|
||||
|
||||
@param row: int
|
||||
@param column: int
|
||||
@param parent: QModelIndex*
|
||||
|
||||
@retval QModelIndex*
|
||||
'''
|
||||
"""
|
||||
if not self.hasIndex(row, column, parent):
|
||||
return QtCore.QModelIndex()
|
||||
|
||||
if not parent.isValid():
|
||||
parent_item = self._root
|
||||
parent_item = self.root_node
|
||||
else:
|
||||
parent_item = parent.internalPointer()
|
||||
|
||||
@@ -152,64 +173,66 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
|
||||
else:
|
||||
return QtCore.QModelIndex()
|
||||
|
||||
def parent(self, mindex):
|
||||
''' get parent of the model item with the given index
|
||||
def parent(self, model_index):
|
||||
""" get parent of the model item with the given index
|
||||
|
||||
if the item has no parent, an invalid QModelIndex* is returned
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
@param model_index: QModelIndex*
|
||||
|
||||
@retval QModelIndex*
|
||||
'''
|
||||
if not mindex.isValid():
|
||||
"""
|
||||
if not model_index.isValid():
|
||||
return QtCore.QModelIndex()
|
||||
|
||||
child = mindex.internalPointer()
|
||||
child = model_index.internalPointer()
|
||||
parent = child.parent()
|
||||
|
||||
if parent == self._root:
|
||||
if parent == self.root_node:
|
||||
return QtCore.QModelIndex()
|
||||
|
||||
return self.createIndex(parent.row(), 0, parent)
|
||||
|
||||
def iterateChildrenIndexFromRootIndex(self, mindex, ignore_root=True):
|
||||
''' depth-first traversal of child nodes
|
||||
def iterateChildrenIndexFromRootIndex(self, model_index, ignore_root=True):
|
||||
""" depth-first traversal of child nodes
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
@param model_index: QModelIndex*
|
||||
@param ignore_root: if set, do not return root index
|
||||
|
||||
@retval yield QModelIndex*
|
||||
'''
|
||||
"""
|
||||
visited = set()
|
||||
stack = deque((mindex,))
|
||||
stack = deque((model_index,))
|
||||
|
||||
while True:
|
||||
try:
|
||||
cmindex = stack.pop()
|
||||
child_index = stack.pop()
|
||||
except IndexError:
|
||||
break
|
||||
|
||||
if cmindex not in visited:
|
||||
if not ignore_root or cmindex is not mindex:
|
||||
if child_index not in visited:
|
||||
if not ignore_root or child_index is not model_index:
|
||||
# ignore root
|
||||
yield cmindex
|
||||
yield child_index
|
||||
|
||||
visited.add(cmindex)
|
||||
visited.add(child_index)
|
||||
|
||||
for idx in range(self.rowCount(cmindex)):
|
||||
stack.append(cmindex.child(idx, 0))
|
||||
for idx in range(self.rowCount(child_index)):
|
||||
stack.append(child_index.child(idx, 0))
|
||||
|
||||
def _util_reset_ida_highlighting(self, item, checked):
|
||||
''' '''
|
||||
if not isinstance(item, (CapaExplorerStringViewItem, CapaExplorerInstructionViewItem, CapaExplorerByteViewItem)):
|
||||
def util_reset_ida_highlighting(self, item, checked):
|
||||
""" """
|
||||
if not isinstance(item, (CapaExplorerStringViewItem, CapaExplorerInstructionViewItem,
|
||||
CapaExplorerByteViewItem)):
|
||||
# ignore other item types
|
||||
return
|
||||
|
||||
curr_highlight = idc.get_color(item.ea, idc.CIC_ITEM)
|
||||
curr_highlight = idc.get_color(item.location, idc.CIC_ITEM)
|
||||
|
||||
if checked:
|
||||
# item checked - record current highlight and set to new
|
||||
item.ida_highlight = curr_highlight
|
||||
idc.set_color(item.ea, idc.CIC_ITEM, DEFAULT_HIGHLIGHT)
|
||||
idc.set_color(item.location, idc.CIC_ITEM, DEFAULT_HIGHLIGHT)
|
||||
else:
|
||||
# item unchecked - reset highlight
|
||||
if curr_highlight != DEFAULT_HIGHLIGHT:
|
||||
@@ -217,36 +240,37 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
|
||||
item.ida_highlight = curr_highlight
|
||||
else:
|
||||
# reset highlight to previous
|
||||
idc.set_color(item.ea, idc.CIC_ITEM, item.ida_highlight)
|
||||
idc.set_color(item.location, idc.CIC_ITEM, item.ida_highlight)
|
||||
|
||||
def setData(self, mindex, value, role):
|
||||
''' set the role data for the item at index to value
|
||||
def setData(self, model_index, value, role):
|
||||
""" set the role data for the item at index to value
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
@param model_index: QModelIndex*
|
||||
@param value: QVariant*
|
||||
@param role: QtCore.Qt.EditRole
|
||||
|
||||
@retval True/False
|
||||
'''
|
||||
if not mindex.isValid():
|
||||
"""
|
||||
if not model_index.isValid():
|
||||
return False
|
||||
|
||||
if role == QtCore.Qt.CheckStateRole and mindex.column() == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION:
|
||||
if role == QtCore.Qt.CheckStateRole and model_index.column() ==\
|
||||
CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION:
|
||||
# user un/checked box - un/check parent and children
|
||||
for cindex in self.iterateChildrenIndexFromRootIndex(mindex, ignore_root=False):
|
||||
cindex.internalPointer().setChecked(value)
|
||||
self._util_reset_ida_highlighting(cindex.internalPointer(), value)
|
||||
self.dataChanged.emit(cindex, cindex)
|
||||
for child_index in self.iterateChildrenIndexFromRootIndex(model_index, ignore_root=False):
|
||||
child_index.internalPointer().setChecked(value)
|
||||
self.util_reset_ida_highlighting(child_index.internalPointer(), value)
|
||||
self.dataChanged.emit(child_index, child_index)
|
||||
return True
|
||||
|
||||
if role == QtCore.Qt.EditRole and value and \
|
||||
mindex.column() == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION and \
|
||||
isinstance(mindex.internalPointer(), CapaExplorerFunctionItem):
|
||||
model_index.column() == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION and \
|
||||
isinstance(model_index.internalPointer(), CapaExplorerFunctionItem):
|
||||
# user renamed function - update IDA database and data model
|
||||
old_name = mindex.internalPointer().info
|
||||
old_name = model_index.internalPointer().info
|
||||
new_name = str(value)
|
||||
|
||||
if idaapi.set_name(mindex.internalPointer().ea, new_name):
|
||||
if idaapi.set_name(model_index.internalPointer().location, new_name):
|
||||
# success update IDA database - update data model
|
||||
self.update_function_name(old_name, new_name)
|
||||
return True
|
||||
@@ -254,167 +278,280 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
|
||||
# no handle
|
||||
return False
|
||||
|
||||
def rowCount(self, mindex):
|
||||
''' get the number of rows under the given parent
|
||||
def rowCount(self, model_index):
|
||||
""" get the number of rows under the given parent
|
||||
|
||||
when the parent is valid it means that is returning the number of
|
||||
children of parent
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
@param model_index: QModelIndex*
|
||||
|
||||
@retval row count
|
||||
'''
|
||||
if mindex.column() > 0:
|
||||
"""
|
||||
if model_index.column() > 0:
|
||||
return 0
|
||||
|
||||
if not mindex.isValid():
|
||||
item = self._root
|
||||
if not model_index.isValid():
|
||||
item = self.root_node
|
||||
else:
|
||||
item = mindex.internalPointer()
|
||||
item = model_index.internalPointer()
|
||||
|
||||
return item.childCount()
|
||||
|
||||
def render_capa_results(self, rule_set, results):
|
||||
''' populate data model with capa results
|
||||
def render_capa_doc_statement_node(self, parent, statement, doc):
|
||||
""" render capa statement read from doc
|
||||
|
||||
@param rule_set: TODO
|
||||
@param results: TODO
|
||||
'''
|
||||
# prepare data model for changes
|
||||
self.beginResetModel()
|
||||
@param parent: parent to which new child is assigned
|
||||
@param statement: statement read from doc
|
||||
@param doc: capa result doc
|
||||
|
||||
for (rule, ress) in results.items():
|
||||
if rule_set.rules[rule].meta.get('lib', False):
|
||||
# skip library rules
|
||||
continue
|
||||
"statement": {
|
||||
"type": "or"
|
||||
},
|
||||
"""
|
||||
if statement['type'] in ('and', 'or', 'optional'):
|
||||
return CapaExplorerDefaultItem(parent, statement['type'])
|
||||
elif statement['type'] == 'not':
|
||||
# TODO: do we display 'not'
|
||||
pass
|
||||
elif statement['type'] == 'some':
|
||||
return CapaExplorerDefaultItem(parent, statement['count'] + ' or more')
|
||||
elif statement['type'] == 'range':
|
||||
# `range` is a weird node, its almost a hybrid of statement + feature.
|
||||
# it is a specific feature repeated multiple times.
|
||||
# there's no additional logic in the feature part, just the existence of a feature.
|
||||
# so, we have to inline some of the feature rendering here.
|
||||
display = 'count(%s): ' % self.capa_doc_feature_to_display(statement['child'])
|
||||
|
||||
# top level item is rule
|
||||
parent = CapaExplorerRuleItem(self._root, rule, len(ress), rule_set.rules[rule].definition)
|
||||
if statement['max'] == statement['min']:
|
||||
display += '%d' % (statement['min'])
|
||||
elif statement['min'] == 0:
|
||||
display += '%d or fewer' % (statement['max'])
|
||||
elif statement['max'] == (1 << 64 - 1):
|
||||
display += '%d or more' % (statement['min'])
|
||||
else:
|
||||
display += 'between %d and %d' % (statement['min'], statement['max'])
|
||||
|
||||
for (ea, res) in sorted(ress, key=lambda p: p[0]):
|
||||
if rule_set.rules[rule].scope == capa.rules.FILE_SCOPE:
|
||||
# file scope - parent is rule
|
||||
parent2 = parent
|
||||
elif rule_set.rules[rule].scope == capa.rules.FUNCTION_SCOPE:
|
||||
parent2 = CapaExplorerFunctionItem(parent, idaapi.get_name(ea), ea)
|
||||
elif rule_set.rules[rule].scope == capa.rules.BASIC_BLOCK_SCOPE:
|
||||
parent2 = CapaExplorerBlockItem(parent, ea)
|
||||
else:
|
||||
# TODO: better way to notify a missed scope?
|
||||
parent2 = CapaExplorerDefaultItem(parent, '', ea)
|
||||
return CapaExplorerFeatureItem(parent, display=display)
|
||||
elif statement['type'] == 'subscope':
|
||||
return CapaExplorerFeatureItem(parent, 'subscope(%s)' % statement['subscope'])
|
||||
elif statement['type'] == 'regex':
|
||||
# regex is a `Statement` not a `Feature`
|
||||
# this is because it doesn't get extracted, but applies to all strings in scope.
|
||||
# so we have to handle it here
|
||||
return CapaExplorerFeatureItem(parent, 'regex(%s)' % statement['pattern'], details=statement['match'])
|
||||
else:
|
||||
raise RuntimeError('unexpected match statement type: ' + str(statement))
|
||||
|
||||
self._render_result(rule_set, res, parent2)
|
||||
def render_capa_doc_match(self, parent, match, doc):
|
||||
""" render capa match read from doc
|
||||
|
||||
# reset data model after making changes
|
||||
self.endResetModel()
|
||||
@param parent: parent node to which new child is assigned
|
||||
@param match: match read from doc
|
||||
@param doc: capa result doc
|
||||
|
||||
def _render_result(self, rule_set, result, parent):
|
||||
''' '''
|
||||
if not result.success:
|
||||
# TODO: display failed branches??
|
||||
"matches": {
|
||||
"0": {
|
||||
"children": [],
|
||||
"locations": [
|
||||
4317184
|
||||
],
|
||||
"node": {
|
||||
"feature": {
|
||||
"section": ".rsrc",
|
||||
"type": "section"
|
||||
},
|
||||
"type": "feature"
|
||||
},
|
||||
"success": true
|
||||
}
|
||||
},
|
||||
"""
|
||||
if not match['success']:
|
||||
# TODO: display failed branches at some point? Help with debugging rules?
|
||||
return
|
||||
|
||||
if isinstance(result.statement, capa.engine.Some):
|
||||
if result.statement.count == 0:
|
||||
if sum(map(lambda c: c.success, result.children)) > 0:
|
||||
parent2 = CapaExplorerDefaultItem(parent, 'optional')
|
||||
else:
|
||||
parent2 = parent
|
||||
else:
|
||||
parent2 = CapaExplorerDefaultItem(parent, '%d or more' % result.statement.count)
|
||||
elif not isinstance(result.statement, (capa.features.Feature, capa.engine.Range, capa.engine.Regex)):
|
||||
# when rending a structural node (and/or/not) then we only care about the node name.
|
||||
if not list(filter(lambda c: bool(c), result.children)):
|
||||
# ignore structural expressions that do not have any successful children (e.g. not)
|
||||
return
|
||||
parent2 = CapaExplorerDefaultItem(parent, result.statement.name.lower())
|
||||
# optional statement with no successful children is empty
|
||||
if (match['node'].get('statement', {}).get('type') == 'optional'
|
||||
and not any(map(lambda m: m['success'], match['children']))):
|
||||
return
|
||||
|
||||
if match['node']['type'] == 'statement':
|
||||
parent2 = self.render_capa_doc_statement_node(parent, match['node']['statement'], doc)
|
||||
elif match['node']['type'] == 'feature':
|
||||
parent2 = self.render_capa_doc_feature_node(parent, match['node']['feature'], match['locations'], doc)
|
||||
else:
|
||||
# but when rendering a Feature, want to see any arguments to it
|
||||
if len(result.locations) == 1:
|
||||
# ea = result.locations.pop()
|
||||
ea = next(iter(result.locations))
|
||||
parent2 = self._render_feature(rule_set, parent, result.statement, ea, str(result.statement))
|
||||
else:
|
||||
parent2 = CapaExplorerDefaultItem(parent, str(result.statement))
|
||||
raise RuntimeError('unexpected node type: ' + str(match['node']['type']))
|
||||
|
||||
for ea in sorted(result.locations):
|
||||
self._render_feature(rule_set, parent2, result.statement, ea)
|
||||
for child in match['children']:
|
||||
self.render_capa_doc_match(parent2, child, doc)
|
||||
|
||||
for child in result.children:
|
||||
self._render_result(rule_set, child, parent2)
|
||||
def render_capa_doc(self, doc):
|
||||
""" render capa features specified in doc
|
||||
|
||||
def _render_feature(self, rule_set, parent, feature, ea, name='-'):
|
||||
''' render a given feature
|
||||
@param doc: capa result doc
|
||||
"""
|
||||
self.beginResetModel()
|
||||
|
||||
for rule in rutils.capability_rules(doc):
|
||||
parent = CapaExplorerRuleItem(self.root_node, rule['meta']['name'], len(rule['matches']), rule['source'])
|
||||
|
||||
for (location, match) in doc[rule['meta']['name']]['matches'].items():
|
||||
if rule['meta']['scope'] == capa.rules.FILE_SCOPE:
|
||||
parent2 = parent
|
||||
elif rule['meta']['scope'] == capa.rules.FUNCTION_SCOPE:
|
||||
parent2 = CapaExplorerFunctionItem(parent, location)
|
||||
elif rule['meta']['scope'] == capa.rules.BASIC_BLOCK_SCOPE:
|
||||
parent2 = CapaExplorerBlockItem(parent, location)
|
||||
else:
|
||||
raise RuntimeError('unexpected rule scope: ' + str(rule['meta']['scope']))
|
||||
|
||||
self.render_capa_doc_match(parent2, match, doc)
|
||||
|
||||
self.endResetModel()
|
||||
|
||||
def capa_doc_feature_to_display(self, feature):
|
||||
""" convert capa doc feature type string to display string for ui
|
||||
|
||||
@param feature: capa feature read from doc
|
||||
|
||||
"feature": {
|
||||
"number": 2147483903,
|
||||
"type": "number"
|
||||
},
|
||||
"""
|
||||
mapping = {
|
||||
'string': 'string(%s)',
|
||||
'bytes': 'bytes(%s)',
|
||||
'api': 'api(%s)',
|
||||
'mnemonic': 'mnemonic(%s)',
|
||||
'export': 'export(%s)',
|
||||
'import': 'import(%s)',
|
||||
'section': 'section(%s)',
|
||||
'number': 'number(0x%X)',
|
||||
'offset': 'offset(0x%X)',
|
||||
'characteristic': 'characteristic(%s)',
|
||||
'match': 'rule match(%s)'
|
||||
}
|
||||
|
||||
@param rule_set: TODO
|
||||
@param parent: TODO
|
||||
@param result: TODO
|
||||
@param ea: virtual address
|
||||
@param name: TODO
|
||||
'''
|
||||
"feature": {
|
||||
"characteristic": [
|
||||
"loop",
|
||||
true
|
||||
],
|
||||
"type": "characteristic"
|
||||
},
|
||||
'''
|
||||
if feature['type'] == 'characteristic':
|
||||
return mapping['characteristic'] % feature['characteristic'][0]
|
||||
|
||||
# convert bytes feature from "410ab4" to "41 0A B4"
|
||||
if feature['type'] == 'bytes':
|
||||
return mapping['bytes'] % ' '.join(feature['bytes'][i:i + 2] for i in
|
||||
range(0, len(feature['bytes']), 2)).upper()
|
||||
|
||||
try:
|
||||
fmt = mapping[feature['type']]
|
||||
except KeyError:
|
||||
raise RuntimeError('unexpected doc type: ' + str(feature['type']))
|
||||
|
||||
return fmt % feature[feature['type']]
|
||||
|
||||
def render_capa_doc_feature_node(self, parent, feature, locations, doc):
|
||||
""" """
|
||||
display = self.capa_doc_feature_to_display(feature)
|
||||
|
||||
if len(locations) == 1:
|
||||
parent2 = self.render_capa_doc_feature(parent, feature, next(iter(locations)), doc, display=display)
|
||||
else:
|
||||
# feature has multiple children, nest under one parent feature node
|
||||
parent2 = CapaExplorerFeatureItem(parent, display)
|
||||
|
||||
for location in sorted(locations):
|
||||
self.render_capa_doc_feature(parent2, feature, location, doc)
|
||||
|
||||
return parent2
|
||||
|
||||
def render_capa_doc_feature(self, parent, feature, location, doc, display='-'):
|
||||
""" render capa feature read from doc
|
||||
|
||||
@param parent: parent node to which new child is assigned
|
||||
@param feature: feature read from doc
|
||||
@param doc: capa feature doc
|
||||
|
||||
"node": {
|
||||
"feature": {
|
||||
"number": 255,
|
||||
"type": "number"
|
||||
},
|
||||
"type": "feature"
|
||||
},
|
||||
|
||||
@param location: address of feature
|
||||
@param display: text to display in plugin ui
|
||||
"""
|
||||
instruction_view = (
|
||||
capa.features.Bytes,
|
||||
capa.features.String,
|
||||
capa.features.insn.API,
|
||||
capa.features.insn.Mnemonic,
|
||||
capa.features.insn.Number,
|
||||
capa.features.insn.Offset
|
||||
'bytes',
|
||||
'api',
|
||||
'mnemonic',
|
||||
'number',
|
||||
'offset'
|
||||
)
|
||||
|
||||
byte_view = (
|
||||
capa.features.file.Section,
|
||||
'section',
|
||||
)
|
||||
|
||||
string_view = (
|
||||
capa.engine.Regex,
|
||||
'string',
|
||||
)
|
||||
default_feature_view = (
|
||||
'import',
|
||||
'export'
|
||||
)
|
||||
|
||||
if isinstance(feature, instruction_view):
|
||||
return CapaExplorerInstructionViewItem(parent, name, ea)
|
||||
# special handling for characteristic pending type
|
||||
if feature['type'] == 'characteristic':
|
||||
if feature['characteristic'][0] in ('embedded pe',):
|
||||
return CapaExplorerByteViewItem(parent, display, location)
|
||||
|
||||
if isinstance(feature, byte_view):
|
||||
return CapaExplorerByteViewItem(parent, name, ea)
|
||||
if feature['characteristic'][0] in ('loop', 'recursive call', 'tight loop', 'switch'):
|
||||
return CapaExplorerFeatureItem(parent, display=display)
|
||||
|
||||
if isinstance(feature, string_view):
|
||||
# TODO: move string collection to item constructor
|
||||
if isinstance(feature, capa.engine.Regex):
|
||||
# rstrip "matched="<string>")" because data already displayed in interface
|
||||
name = name.split(',')[0] + ')'
|
||||
return CapaExplorerStringViewItem(parent, name, ea, feature.match)
|
||||
# default to instruction view
|
||||
return CapaExplorerInstructionViewItem(parent, display, location)
|
||||
|
||||
if isinstance(feature, capa.features.Characteristic):
|
||||
# special rendering for characteristics
|
||||
if feature.name in ('loop', 'recursive call', 'tight loop', 'switch'):
|
||||
return CapaExplorerDefaultItem(parent, name)
|
||||
if feature.name in ('embedded pe',):
|
||||
return CapaExplorerByteViewItem(parent, name, ea)
|
||||
return CapaExplorerInstructionViewItem(parent, name, ea)
|
||||
if feature['type'] == 'match':
|
||||
return CapaExplorerRuleMatchItem(parent, display, source=doc.get(feature['match'], {}).get('source', ''))
|
||||
|
||||
if isinstance(feature, capa.features.MatchedRule):
|
||||
# render feature as a rule item
|
||||
return CapaExplorerRuleItem(parent, name, 0, rule_set.rules[feature.rule_name].definition)
|
||||
if feature['type'] in instruction_view:
|
||||
return CapaExplorerInstructionViewItem(parent, display, location)
|
||||
|
||||
if isinstance(feature, capa.engine.Range):
|
||||
# render feature based upon type child
|
||||
return self._render_feature(rule_set, parent, feature.child, ea, name)
|
||||
if feature['type'] in byte_view:
|
||||
return CapaExplorerByteViewItem(parent, display, location)
|
||||
|
||||
# no handle, default to name and virtual address display
|
||||
return CapaExplorerDefaultItem(parent, name, ea)
|
||||
if feature['type'] in string_view:
|
||||
return CapaExplorerStringViewItem(parent, display, location)
|
||||
|
||||
if feature['type'] in default_feature_view:
|
||||
return CapaExplorerFeatureItem(parent, display=display)
|
||||
|
||||
raise RuntimeError('unexpected feature type: ' + str(feature['type']))
|
||||
|
||||
def update_function_name(self, old_name, new_name):
|
||||
''' update all instances of function name
|
||||
""" update all instances of function name
|
||||
|
||||
@param old_name: previous function name
|
||||
@param new_name: new function name
|
||||
'''
|
||||
rmindex = self.index(0, 0, QtCore.QModelIndex())
|
||||
"""
|
||||
root_index = self.index(0, 0, QtCore.QModelIndex())
|
||||
|
||||
# convert name to view format for matching
|
||||
# TODO: handle this better
|
||||
old_name = CapaExplorerFunctionItem.view_fmt % old_name
|
||||
old_name = CapaExplorerFunctionItem.fmt % old_name
|
||||
|
||||
for mindex in self.match(rmindex, QtCore.Qt.DisplayRole, old_name, hits=-1, flags=QtCore.Qt.MatchRecursive):
|
||||
if not isinstance(mindex.internalPointer(), CapaExplorerFunctionItem):
|
||||
for model_index in self.match(root_index, QtCore.Qt.DisplayRole, old_name, hits=-1,
|
||||
flags=QtCore.Qt.MatchRecursive):
|
||||
if not isinstance(model_index.internalPointer(), CapaExplorerFunctionItem):
|
||||
continue
|
||||
mindex.internalPointer().info = new_name
|
||||
self.dataChanged.emit(mindex, mindex)
|
||||
|
||||
model_index.internalPointer().info = new_name
|
||||
self.dataChanged.emit(model_index, model_index)
|
||||
|
||||
@@ -1,25 +1,27 @@
|
||||
from PyQt5 import QtCore
|
||||
|
||||
from capa.ida.explorer.model import CapaExplorerDataModel
|
||||
|
||||
|
||||
class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel):
|
||||
|
||||
def __init__(self, parent=None):
|
||||
''' '''
|
||||
""" """
|
||||
super(CapaExplorerSortFilterProxyModel, self).__init__(parent)
|
||||
|
||||
def lessThan(self, left, right):
|
||||
''' true if the value of the left item is less than value of right item
|
||||
""" true if the value of the left item is less than value of right item
|
||||
|
||||
@param left: QModelIndex*
|
||||
@param right: QModelIndex*
|
||||
|
||||
@retval True/False
|
||||
'''
|
||||
"""
|
||||
ldata = left.internalPointer().data(left.column())
|
||||
rdata = right.internalPointer().data(right.column())
|
||||
|
||||
if ldata and rdata and left.column() == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS and left.column() == right.column():
|
||||
if ldata and rdata and left.column() == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS and left.column() \
|
||||
== right.column():
|
||||
# convert virtual address before compare
|
||||
return int(ldata, 16) < int(rdata, 16)
|
||||
else:
|
||||
@@ -27,49 +29,49 @@ class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel):
|
||||
return ldata.lower() < rdata.lower()
|
||||
|
||||
def filterAcceptsRow(self, row, parent):
|
||||
''' true if the item in the row indicated by the given row and parent
|
||||
""" true if the item in the row indicated by the given row and parent
|
||||
should be included in the model; otherwise returns false
|
||||
@param row: int
|
||||
@param parent: QModelIndex*
|
||||
|
||||
@retval True/False
|
||||
'''
|
||||
if self._filter_accepts_row_self(row, parent):
|
||||
"""
|
||||
if self.filter_accepts_row_self(row, parent):
|
||||
return True
|
||||
|
||||
alpha = parent
|
||||
while alpha.isValid():
|
||||
if self._filter_accepts_row_self(alpha.row(), alpha.parent()):
|
||||
if self.filter_accepts_row_self(alpha.row(), alpha.parent()):
|
||||
return True
|
||||
alpha = alpha.parent()
|
||||
|
||||
if self._index_has_accepted_children(row, parent):
|
||||
if self.index_has_accepted_children(row, parent):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def add_single_string_filter(self, column, string):
|
||||
''' add fixed string filter
|
||||
""" add fixed string filter
|
||||
|
||||
@param column: key column
|
||||
@param string: string to sort
|
||||
'''
|
||||
"""
|
||||
self.setFilterKeyColumn(column)
|
||||
self.setFilterFixedString(string)
|
||||
|
||||
def _index_has_accepted_children(self, row, parent):
|
||||
''' '''
|
||||
mindex = self.sourceModel().index(row, 0, parent)
|
||||
def index_has_accepted_children(self, row, parent):
|
||||
""" """
|
||||
model_index = self.sourceModel().index(row, 0, parent)
|
||||
|
||||
if mindex.isValid():
|
||||
for idx in range(self.sourceModel().rowCount(mindex)):
|
||||
if self._filter_accepts_row_self(idx, mindex):
|
||||
if model_index.isValid():
|
||||
for idx in range(self.sourceModel().rowCount(model_index)):
|
||||
if self.filter_accepts_row_self(idx, model_index):
|
||||
return True
|
||||
if self._index_has_accepted_children(idx, mindex):
|
||||
if self.index_has_accepted_children(idx, model_index):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _filter_accepts_row_self(self, row, parent):
|
||||
''' '''
|
||||
def filter_accepts_row_self(self, row, parent):
|
||||
""" """
|
||||
return super(CapaExplorerSortFilterProxyModel, self).filterAcceptsRow(row, parent)
|
||||
|
||||
@@ -4,40 +4,43 @@ import idaapi
|
||||
import idc
|
||||
|
||||
from capa.ida.explorer.model import CapaExplorerDataModel
|
||||
from capa.ida.explorer.item import CapaExplorerFunctionItem
|
||||
from capa.ida.explorer.item import (
|
||||
CapaExplorerFunctionItem,
|
||||
CapaExplorerRuleItem,
|
||||
)
|
||||
|
||||
|
||||
class CapaExplorerQtreeView(QtWidgets.QTreeView):
|
||||
''' capa explorer QTreeView implementation
|
||||
""" capa explorer QTreeView implementation
|
||||
|
||||
view controls UI action responses and displays data from
|
||||
CapaExplorerDataModel
|
||||
|
||||
view does not modify CapaExplorerDataModel directly - data
|
||||
modifications should be implemented in CapaExplorerDataModel
|
||||
'''
|
||||
"""
|
||||
|
||||
def __init__(self, model, parent=None):
|
||||
''' initialize CapaExplorerQTreeView
|
||||
""" initialize CapaExplorerQTreeView
|
||||
|
||||
TODO
|
||||
|
||||
@param model: TODO
|
||||
@param parent: TODO
|
||||
'''
|
||||
"""
|
||||
super(CapaExplorerQtreeView, self).__init__(parent)
|
||||
|
||||
self.setModel(model)
|
||||
|
||||
# TODO: get from parent??
|
||||
self._model = model
|
||||
self._parent = parent
|
||||
self.model = model
|
||||
self.parent = parent
|
||||
|
||||
# configure custom UI controls
|
||||
self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
|
||||
self.setExpandsOnDoubleClick(False)
|
||||
self.setSortingEnabled(True)
|
||||
self._model.setDynamicSortFilter(False)
|
||||
self.model.setDynamicSortFilter(False)
|
||||
|
||||
# configure view columns to auto-resize
|
||||
for idx in range(CapaExplorerDataModel.COLUMN_COUNT):
|
||||
@@ -48,223 +51,222 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView):
|
||||
self.collapsed.connect(self.resize_columns_to_content)
|
||||
|
||||
# connect slots
|
||||
self.customContextMenuRequested.connect(self._slot_custom_context_menu_requested)
|
||||
self.doubleClicked.connect(self._slot_double_click)
|
||||
# self.clicked.connect(self._slot_click)
|
||||
self.customContextMenuRequested.connect(self.slot_custom_context_menu_requested)
|
||||
self.doubleClicked.connect(self.slot_double_click)
|
||||
# self.clicked.connect(self.slot_click)
|
||||
|
||||
self.setStyleSheet('QTreeView::item {padding-right: 15 px;padding-bottom: 2 px;}')
|
||||
|
||||
def reset(self):
|
||||
''' reset user interface changes
|
||||
""" reset user interface changes
|
||||
|
||||
called when view should reset any user interface changes
|
||||
made since the last reset e.g. IDA window highlighting
|
||||
'''
|
||||
"""
|
||||
self.collapseAll()
|
||||
self.resize_columns_to_content()
|
||||
|
||||
def resize_columns_to_content(self):
|
||||
''' reset view columns to contents
|
||||
""" reset view columns to contents
|
||||
|
||||
TODO: prevent columns from shrinking
|
||||
'''
|
||||
"""
|
||||
self.header().resizeSections(QtWidgets.QHeaderView.ResizeToContents)
|
||||
|
||||
def _map_index_to_source_item(self, mindex):
|
||||
''' map proxy model index to source model item
|
||||
def map_index_to_source_item(self, model_index):
|
||||
""" map proxy model index to source model item
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
@param model_index: QModelIndex*
|
||||
|
||||
@retval QObject*
|
||||
'''
|
||||
return self._model.mapToSource(mindex).internalPointer()
|
||||
"""
|
||||
return self.model.mapToSource(model_index).internalPointer()
|
||||
|
||||
def _send_data_to_clipboard(self, data):
|
||||
''' copy data to the clipboard
|
||||
def send_data_to_clipboard(self, data):
|
||||
""" copy data to the clipboard
|
||||
|
||||
@param data: data to be copied
|
||||
'''
|
||||
"""
|
||||
clip = QtWidgets.QApplication.clipboard()
|
||||
clip.clear(mode=clip.Clipboard)
|
||||
clip.setText(data, mode=clip.Clipboard)
|
||||
|
||||
def _new_action(self, display, data, slot):
|
||||
''' create action for context menu
|
||||
def new_action(self, display, data, slot):
|
||||
""" create action for context menu
|
||||
|
||||
@param display: text displayed to user in context menu
|
||||
@param data: data passed to slot
|
||||
@param slot: slot to connect
|
||||
|
||||
@retval QAction*
|
||||
'''
|
||||
action = QtWidgets.QAction(display, self._parent)
|
||||
"""
|
||||
action = QtWidgets.QAction(display, self.parent)
|
||||
action.setData(data)
|
||||
action.triggered.connect(lambda checked: slot(action))
|
||||
|
||||
return action
|
||||
|
||||
def _load_default_context_menu_actions(self, data):
|
||||
''' yield actions specific to function custom context menu
|
||||
def load_default_context_menu_actions(self, data):
|
||||
""" yield actions specific to function custom context menu
|
||||
|
||||
@param data: tuple
|
||||
|
||||
@yield QAction*
|
||||
'''
|
||||
"""
|
||||
default_actions = [
|
||||
('Copy column', data, self._slot_copy_column),
|
||||
('Copy row', data, self._slot_copy_row),
|
||||
# ('Filter', data, self._slot_filter),
|
||||
('Copy column', data, self.slot_copy_column),
|
||||
('Copy row', data, self.slot_copy_row),
|
||||
]
|
||||
|
||||
# add default actions
|
||||
for action in default_actions:
|
||||
yield self._new_action(*action)
|
||||
yield self.new_action(*action)
|
||||
|
||||
def _load_function_context_menu_actions(self, data):
|
||||
''' yield actions specific to function custom context menu
|
||||
def load_function_context_menu_actions(self, data):
|
||||
""" yield actions specific to function custom context menu
|
||||
|
||||
@param data: tuple
|
||||
|
||||
@yield QAction*
|
||||
'''
|
||||
"""
|
||||
function_actions = [
|
||||
('Rename function', data, self._slot_rename_function),
|
||||
('Rename function', data, self.slot_rename_function),
|
||||
]
|
||||
|
||||
# add function actions
|
||||
for action in function_actions:
|
||||
yield self._new_action(*action)
|
||||
yield self.new_action(*action)
|
||||
|
||||
# add default actions
|
||||
for action in self._load_default_context_menu_actions(data):
|
||||
for action in self.load_default_context_menu_actions(data):
|
||||
yield action
|
||||
|
||||
def _load_default_context_menu(self, pos, item, mindex):
|
||||
''' create default custom context menu
|
||||
def load_default_context_menu(self, pos, item, model_index):
|
||||
""" create default custom context menu
|
||||
|
||||
creates custom context menu containing default actions
|
||||
|
||||
@param pos: TODO
|
||||
@param item: TODO
|
||||
@param mindex: TODO
|
||||
@param model_index: TODO
|
||||
|
||||
@retval QMenu*
|
||||
'''
|
||||
"""
|
||||
menu = QtWidgets.QMenu()
|
||||
|
||||
for action in self._load_default_context_menu_actions((pos, item, mindex)):
|
||||
for action in self.load_default_context_menu_actions((pos, item, model_index)):
|
||||
menu.addAction(action)
|
||||
|
||||
return menu
|
||||
|
||||
def _load_function_item_context_menu(self, pos, item, mindex):
|
||||
''' create function custom context menu
|
||||
def load_function_item_context_menu(self, pos, item, model_index):
|
||||
""" create function custom context menu
|
||||
|
||||
creates custom context menu containing actions specific to functions
|
||||
and the default actions
|
||||
|
||||
@param pos: TODO
|
||||
@param item: TODO
|
||||
@param mindex: TODO
|
||||
@param model_index: TODO
|
||||
|
||||
@retval QMenu*
|
||||
'''
|
||||
"""
|
||||
menu = QtWidgets.QMenu()
|
||||
|
||||
for action in self._load_function_context_menu_actions((pos, item, mindex)):
|
||||
for action in self.load_function_context_menu_actions((pos, item, model_index)):
|
||||
menu.addAction(action)
|
||||
|
||||
return menu
|
||||
|
||||
def _show_custom_context_menu(self, menu, pos):
|
||||
''' display custom context menu in view
|
||||
def show_custom_context_menu(self, menu, pos):
|
||||
""" display custom context menu in view
|
||||
|
||||
@param menu: TODO
|
||||
@param pos: TODO
|
||||
'''
|
||||
"""
|
||||
if not menu:
|
||||
return
|
||||
|
||||
menu.exec_(self.viewport().mapToGlobal(pos))
|
||||
|
||||
def _slot_copy_column(self, action):
|
||||
''' slot connected to custom context menu
|
||||
def slot_copy_column(self, action):
|
||||
""" slot connected to custom context menu
|
||||
|
||||
allows user to select a column and copy the data
|
||||
to clipboard
|
||||
|
||||
@param action: QAction*
|
||||
'''
|
||||
_, item, mindex = action.data()
|
||||
self._send_data_to_clipboard(item.data(mindex.column()))
|
||||
"""
|
||||
_, item, model_index = action.data()
|
||||
self.send_data_to_clipboard(item.data(model_index.column()))
|
||||
|
||||
def _slot_copy_row(self, action):
|
||||
''' slot connected to custom context menu
|
||||
def slot_copy_row(self, action):
|
||||
""" slot connected to custom context menu
|
||||
|
||||
allows user to select a row and copy the space-delimeted
|
||||
data to clipboard
|
||||
|
||||
@param action: QAction*
|
||||
'''
|
||||
"""
|
||||
_, item, _ = action.data()
|
||||
self._send_data_to_clipboard(str(item))
|
||||
self.send_data_to_clipboard(str(item))
|
||||
|
||||
def _slot_rename_function(self, action):
|
||||
''' slot connected to custom context menu
|
||||
def slot_rename_function(self, action):
|
||||
""" slot connected to custom context menu
|
||||
|
||||
allows user to select a edit a function name and push
|
||||
changes to IDA
|
||||
|
||||
@param action: QAction*
|
||||
'''
|
||||
_, item, mindex = action.data()
|
||||
"""
|
||||
_, item, model_index = action.data()
|
||||
|
||||
# make item temporary edit, reset after user is finished
|
||||
item.setIsEditable(True)
|
||||
self.edit(mindex)
|
||||
self.edit(model_index)
|
||||
item.setIsEditable(False)
|
||||
|
||||
def _slot_custom_context_menu_requested(self, pos):
|
||||
''' slot connected to custom context menu request
|
||||
def slot_custom_context_menu_requested(self, pos):
|
||||
""" slot connected to custom context menu request
|
||||
|
||||
displays custom context menu to user containing action
|
||||
relevant to the data item selected
|
||||
|
||||
@param pos: TODO
|
||||
'''
|
||||
mindex = self.indexAt(pos)
|
||||
"""
|
||||
model_index = self.indexAt(pos)
|
||||
|
||||
if not mindex.isValid():
|
||||
if not model_index.isValid():
|
||||
return
|
||||
|
||||
item = self._map_index_to_source_item(mindex)
|
||||
column = mindex.column()
|
||||
item = self.map_index_to_source_item(model_index)
|
||||
column = model_index.column()
|
||||
menu = None
|
||||
|
||||
if CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == column and isinstance(item, CapaExplorerFunctionItem):
|
||||
# user hovered function item
|
||||
menu = self._load_function_item_context_menu(pos, item, mindex)
|
||||
menu = self.load_function_item_context_menu(pos, item, model_index)
|
||||
else:
|
||||
# user hovered default item
|
||||
menu = self._load_default_context_menu(pos, item, mindex)
|
||||
menu = self.load_default_context_menu(pos, item, model_index)
|
||||
|
||||
# show custom context menu at view position
|
||||
self._show_custom_context_menu(menu, pos)
|
||||
self.show_custom_context_menu(menu, pos)
|
||||
|
||||
def _slot_click(self):
|
||||
''' slot connected to single click event '''
|
||||
def slot_click(self):
|
||||
""" slot connected to single click event """
|
||||
pass
|
||||
|
||||
def _slot_double_click(self, mindex):
|
||||
''' slot connected to double click event
|
||||
def slot_double_click(self, model_index):
|
||||
""" slot connected to double click event
|
||||
|
||||
@param mindex: QModelIndex*
|
||||
'''
|
||||
if not mindex.isValid():
|
||||
@param model_index: QModelIndex*
|
||||
"""
|
||||
if not model_index.isValid():
|
||||
return
|
||||
|
||||
item = self._map_index_to_source_item(mindex)
|
||||
column = mindex.column()
|
||||
item = self.map_index_to_source_item(model_index)
|
||||
column = model_index.column()
|
||||
|
||||
if CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS == column:
|
||||
# user double-clicked virtual address column - navigate IDA to address
|
||||
@@ -275,7 +277,7 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView):
|
||||
|
||||
if CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == column:
|
||||
# user double-clicked information column - un/expand
|
||||
if self.isExpanded(mindex):
|
||||
self.collapse(mindex)
|
||||
if self.isExpanded(model_index):
|
||||
self.collapse(model_index)
|
||||
else:
|
||||
self.expand(mindex)
|
||||
self.expand(model_index)
|
||||
|
||||
@@ -2,24 +2,11 @@ import os
|
||||
import logging
|
||||
import collections
|
||||
|
||||
from PyQt5.QtWidgets import (
|
||||
QHeaderView,
|
||||
QAbstractItemView,
|
||||
QMenuBar,
|
||||
QAction,
|
||||
QTabWidget,
|
||||
QWidget,
|
||||
QTextEdit,
|
||||
QMenu,
|
||||
QApplication,
|
||||
QVBoxLayout,
|
||||
QToolTip,
|
||||
QCheckBox,
|
||||
QTableWidget,
|
||||
QTableWidgetItem
|
||||
from PyQt5 import (
|
||||
QtWidgets,
|
||||
QtGui,
|
||||
QtCore
|
||||
)
|
||||
from PyQt5.QtGui import QCursor, QIcon
|
||||
from PyQt5.QtCore import Qt
|
||||
|
||||
import idaapi
|
||||
|
||||
@@ -27,6 +14,7 @@ import capa.main
|
||||
import capa.rules
|
||||
import capa.features.extractors.ida
|
||||
import capa.ida.helpers
|
||||
import capa.render.utils as rutils
|
||||
|
||||
from capa.ida.explorer.view import CapaExplorerQtreeView
|
||||
from capa.ida.explorer.model import CapaExplorerDataModel
|
||||
@@ -40,254 +28,286 @@ logger = logging.getLogger('capa')
|
||||
class CapaExplorerIdaHooks(idaapi.UI_Hooks):
|
||||
|
||||
def __init__(self, screen_ea_changed_hook, action_hooks):
|
||||
''' facilitate IDA UI hooks
|
||||
""" facilitate IDA UI hooks
|
||||
|
||||
@param screen_ea_changed: TODO
|
||||
@param action_hooks: TODO
|
||||
'''
|
||||
"""
|
||||
super(CapaExplorerIdaHooks, self).__init__()
|
||||
|
||||
self._screen_ea_changed_hook = screen_ea_changed_hook
|
||||
self._process_action_hooks = action_hooks
|
||||
self._process_action_handle = None
|
||||
self._process_action_meta = {}
|
||||
self.screen_ea_changed_hook = screen_ea_changed_hook
|
||||
self.process_action_hooks = action_hooks
|
||||
self.process_action_handle = None
|
||||
self.process_action_meta = {}
|
||||
|
||||
def preprocess_action(self, name):
|
||||
''' called prior to action completed
|
||||
""" called prior to action completed
|
||||
|
||||
@param name: name of action defined by idagui.cfg
|
||||
|
||||
@retval must be 0
|
||||
'''
|
||||
self._process_action_handle = self._process_action_hooks.get(name, None)
|
||||
"""
|
||||
self.process_action_handle = self.process_action_hooks.get(name, None)
|
||||
|
||||
if self._process_action_handle:
|
||||
self._process_action_handle(self._process_action_meta)
|
||||
if self.process_action_handle:
|
||||
self.process_action_handle(self.process_action_meta)
|
||||
|
||||
# must return 0 for IDA
|
||||
return 0
|
||||
|
||||
def postprocess_action(self):
|
||||
''' called after action completed '''
|
||||
if not self._process_action_handle:
|
||||
""" called after action completed """
|
||||
if not self.process_action_handle:
|
||||
return
|
||||
|
||||
self._process_action_handle(self._process_action_meta, post=True)
|
||||
self._reset()
|
||||
self.process_action_handle(self.process_action_meta, post=True)
|
||||
self.reset()
|
||||
|
||||
def screen_ea_changed(self, curr_ea, prev_ea):
|
||||
''' called after screen ea is changed
|
||||
""" called after screen location is changed
|
||||
|
||||
@param curr_ea: current ea
|
||||
@param prev_ea: prev ea
|
||||
'''
|
||||
self._screen_ea_changed_hook(idaapi.get_current_widget(), curr_ea, prev_ea)
|
||||
@param curr_ea: current location
|
||||
@param prev_ea: prev location
|
||||
"""
|
||||
self.screen_ea_changed_hook(idaapi.get_current_widget(), curr_ea, prev_ea)
|
||||
|
||||
def _reset(self):
|
||||
''' reset internal state '''
|
||||
self._process_action_handle = None
|
||||
self._process_action_meta.clear()
|
||||
def reset(self):
|
||||
""" reset internal state """
|
||||
self.process_action_handle = None
|
||||
self.process_action_meta.clear()
|
||||
|
||||
|
||||
class CapaExplorerForm(idaapi.PluginForm):
|
||||
|
||||
def __init__(self):
|
||||
''' '''
|
||||
""" """
|
||||
super(CapaExplorerForm, self).__init__()
|
||||
|
||||
self.form_title = PLUGIN_NAME
|
||||
self.parent = None
|
||||
self._file_loc = __file__
|
||||
self._ida_hooks = None
|
||||
self.file_loc = __file__
|
||||
self.ida_hooks = None
|
||||
|
||||
# models
|
||||
self._model_data = None
|
||||
self._model_proxy = None
|
||||
self.model_data = None
|
||||
self.model_proxy = None
|
||||
|
||||
# user interface elements
|
||||
self._view_limit_results_by_function = None
|
||||
self._view_tree = None
|
||||
self._view_summary = None
|
||||
self._view_tabs = None
|
||||
self._view_menu_bar = None
|
||||
self.view_limit_results_by_function = None
|
||||
self.view_tree = None
|
||||
self.view_summary = None
|
||||
self.view_attack = None
|
||||
self.view_tabs = None
|
||||
self.view_menu_bar = None
|
||||
|
||||
def OnCreate(self, form):
|
||||
''' '''
|
||||
""" """
|
||||
self.parent = self.FormToPyQtWidget(form)
|
||||
self._load_interface()
|
||||
self._load_capa_results()
|
||||
self._load_ida_hooks()
|
||||
self.load_interface()
|
||||
self.load_capa_results()
|
||||
self.load_ida_hooks()
|
||||
|
||||
self._view_tree.reset()
|
||||
self.view_tree.reset()
|
||||
|
||||
logger.info('form created.')
|
||||
|
||||
def Show(self):
|
||||
''' '''
|
||||
""" """
|
||||
return idaapi.PluginForm.Show(self, self.form_title, options=(
|
||||
idaapi.PluginForm.WOPN_TAB | idaapi.PluginForm.WCLS_CLOSE_LATER
|
||||
))
|
||||
|
||||
def OnClose(self, form):
|
||||
''' form is closed '''
|
||||
self._unload_ida_hooks()
|
||||
self._ida_reset()
|
||||
""" form is closed """
|
||||
self.unload_ida_hooks()
|
||||
self.ida_reset()
|
||||
|
||||
logger.info('form closed.')
|
||||
|
||||
def _load_interface(self):
|
||||
''' load user interface '''
|
||||
def load_interface(self):
|
||||
""" load user interface """
|
||||
# load models
|
||||
self._model_data = CapaExplorerDataModel()
|
||||
self._model_proxy = CapaExplorerSortFilterProxyModel()
|
||||
self._model_proxy.setSourceModel(self._model_data)
|
||||
self.model_data = CapaExplorerDataModel()
|
||||
self.model_proxy = CapaExplorerSortFilterProxyModel()
|
||||
self.model_proxy.setSourceModel(self.model_data)
|
||||
|
||||
# load tree
|
||||
self._view_tree = CapaExplorerQtreeView(self._model_proxy, self.parent)
|
||||
self.view_tree = CapaExplorerQtreeView(self.model_proxy, self.parent)
|
||||
|
||||
# load summary table
|
||||
self._load_view_summary()
|
||||
self.load_view_summary()
|
||||
self.load_view_attack()
|
||||
|
||||
# load parent tab and children tab views
|
||||
self._load_view_tabs()
|
||||
self._load_view_checkbox_limit_by()
|
||||
self._load_view_summary_tab()
|
||||
self._load_view_tree_tab()
|
||||
self.load_view_tabs()
|
||||
self.load_view_checkbox_limit_by()
|
||||
self.load_view_summary_tab()
|
||||
self.load_view_attack_tab()
|
||||
self.load_view_tree_tab()
|
||||
|
||||
# load menu bar and sub menus
|
||||
self._load_view_menu_bar()
|
||||
self._load_file_menu()
|
||||
self.load_view_menu_bar()
|
||||
self.load_file_menu()
|
||||
|
||||
# load parent view
|
||||
self._load_view_parent()
|
||||
self.load_view_parent()
|
||||
|
||||
def _load_view_tabs(self):
|
||||
''' '''
|
||||
tabs = QTabWidget()
|
||||
def load_view_tabs(self):
|
||||
""" """
|
||||
tabs = QtWidgets.QTabWidget()
|
||||
self.view_tabs = tabs
|
||||
|
||||
self._view_tabs = tabs
|
||||
def load_view_menu_bar(self):
|
||||
""" """
|
||||
bar = QtWidgets.QMenuBar()
|
||||
self.view_menu_bar = bar
|
||||
|
||||
def _load_view_menu_bar(self):
|
||||
''' '''
|
||||
bar = QMenuBar()
|
||||
# bar.hovered.connect(self._slot_menu_bar_hovered)
|
||||
def load_view_summary(self):
|
||||
""" """
|
||||
table_headers = [
|
||||
'Capability',
|
||||
'Namespace',
|
||||
]
|
||||
|
||||
self._view_menu_bar = bar
|
||||
table = QtWidgets.QTableWidget()
|
||||
|
||||
def _load_view_summary(self):
|
||||
''' '''
|
||||
table = QTableWidget()
|
||||
|
||||
table.setColumnCount(4)
|
||||
table.setColumnCount(len(table_headers))
|
||||
table.verticalHeader().setVisible(False)
|
||||
table.setSortingEnabled(False)
|
||||
table.setEditTriggers(QAbstractItemView.NoEditTriggers)
|
||||
table.setFocusPolicy(Qt.NoFocus)
|
||||
table.setSelectionMode(QAbstractItemView.NoSelection)
|
||||
table.setHorizontalHeaderLabels([
|
||||
'Objectives',
|
||||
'Behaviors',
|
||||
'Techniques',
|
||||
'Rule Hits'
|
||||
])
|
||||
table.horizontalHeader().setDefaultAlignment(Qt.AlignLeft)
|
||||
table.setStyleSheet('QTableWidget::item { border: none; padding: 15px; }')
|
||||
table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
|
||||
table.setFocusPolicy(QtCore.Qt.NoFocus)
|
||||
table.setSelectionMode(QtWidgets.QAbstractItemView.NoSelection)
|
||||
table.setHorizontalHeaderLabels(table_headers)
|
||||
table.horizontalHeader().setDefaultAlignment(QtCore.Qt.AlignLeft)
|
||||
table.setShowGrid(False)
|
||||
table.setStyleSheet('QTableWidget::item { padding: 25px; }')
|
||||
|
||||
self._view_summary = table
|
||||
self.view_summary = table
|
||||
|
||||
def _load_view_checkbox_limit_by(self):
|
||||
''' '''
|
||||
check = QCheckBox('Limit results to current function')
|
||||
def load_view_attack(self):
|
||||
""" """
|
||||
table_headers = [
|
||||
'ATT&CK Tactic',
|
||||
'ATT&CK Technique ',
|
||||
]
|
||||
|
||||
table = QtWidgets.QTableWidget()
|
||||
|
||||
table.setColumnCount(len(table_headers))
|
||||
table.verticalHeader().setVisible(False)
|
||||
table.setSortingEnabled(False)
|
||||
table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
|
||||
table.setFocusPolicy(QtCore.Qt.NoFocus)
|
||||
table.setSelectionMode(QtWidgets.QAbstractItemView.NoSelection)
|
||||
table.setHorizontalHeaderLabels(table_headers)
|
||||
table.horizontalHeader().setDefaultAlignment(QtCore.Qt.AlignLeft)
|
||||
table.setShowGrid(False)
|
||||
table.setStyleSheet('QTableWidget::item { padding: 25px; }')
|
||||
|
||||
self.view_attack = table
|
||||
|
||||
def load_view_checkbox_limit_by(self):
|
||||
""" """
|
||||
check = QtWidgets.QCheckBox('Limit results to current function')
|
||||
check.setChecked(False)
|
||||
check.stateChanged.connect(self._slot_checkbox_limit_by_changed)
|
||||
check.stateChanged.connect(self.slot_checkbox_limit_by_changed)
|
||||
|
||||
self._view_checkbox_limit_by = check
|
||||
self.view_checkbox_limit_by = check
|
||||
|
||||
def _load_view_parent(self):
|
||||
''' load view parent '''
|
||||
layout = QVBoxLayout()
|
||||
layout.addWidget(self._view_tabs)
|
||||
layout.setMenuBar(self._view_menu_bar)
|
||||
def load_view_parent(self):
|
||||
""" load view parent """
|
||||
layout = QtWidgets.QVBoxLayout()
|
||||
layout.addWidget(self.view_tabs)
|
||||
layout.setMenuBar(self.view_menu_bar)
|
||||
|
||||
self.parent.setLayout(layout)
|
||||
|
||||
def _load_view_tree_tab(self):
|
||||
''' load view tree tab '''
|
||||
layout = QVBoxLayout()
|
||||
layout.addWidget(self._view_checkbox_limit_by)
|
||||
layout.addWidget(self._view_tree)
|
||||
def load_view_tree_tab(self):
|
||||
""" load view tree tab """
|
||||
layout = QtWidgets.QVBoxLayout()
|
||||
layout.addWidget(self.view_checkbox_limit_by)
|
||||
layout.addWidget(self.view_tree)
|
||||
|
||||
tab = QWidget()
|
||||
tab = QtWidgets.QWidget()
|
||||
tab.setLayout(layout)
|
||||
|
||||
self._view_tabs.addTab(tab, 'Tree View')
|
||||
self.view_tabs.addTab(tab, 'Tree View')
|
||||
|
||||
def _load_view_summary_tab(self):
|
||||
''' '''
|
||||
layout = QVBoxLayout()
|
||||
layout.addWidget(self._view_summary)
|
||||
def load_view_summary_tab(self):
|
||||
""" """
|
||||
layout = QtWidgets.QVBoxLayout()
|
||||
layout.addWidget(self.view_summary)
|
||||
|
||||
tab = QWidget()
|
||||
tab = QtWidgets.QWidget()
|
||||
tab.setLayout(layout)
|
||||
|
||||
self._view_tabs.addTab(tab, 'Summary')
|
||||
self.view_tabs.addTab(tab, 'Summary')
|
||||
|
||||
def _load_file_menu(self):
|
||||
''' load file menu actions '''
|
||||
def load_view_attack_tab(self):
|
||||
""" """
|
||||
layout = QtWidgets.QVBoxLayout()
|
||||
layout.addWidget(self.view_attack)
|
||||
|
||||
tab = QtWidgets.QWidget()
|
||||
tab.setLayout(layout)
|
||||
|
||||
self.view_tabs.addTab(tab, 'MITRE')
|
||||
|
||||
def load_file_menu(self):
|
||||
""" load file menu actions """
|
||||
actions = (
|
||||
('Reset view', 'Reset plugin view', self.reset),
|
||||
('Run analysis', 'Run capa analysis on current database', self.reload),
|
||||
)
|
||||
|
||||
menu = self._view_menu_bar.addMenu('File')
|
||||
menu = self.view_menu_bar.addMenu('File')
|
||||
|
||||
for name, _, handle in actions:
|
||||
action = QAction(name, self.parent)
|
||||
action = QtWidgets.QAction(name, self.parent)
|
||||
action.triggered.connect(handle)
|
||||
# action.setToolTip(tip)
|
||||
menu.addAction(action)
|
||||
|
||||
def _load_ida_hooks(self):
|
||||
''' '''
|
||||
def load_ida_hooks(self):
|
||||
""" """
|
||||
action_hooks = {
|
||||
'MakeName': self._ida_hook_rename,
|
||||
'EditFunction': self._ida_hook_rename,
|
||||
'MakeName': self.ida_hook_rename,
|
||||
'EditFunction': self.ida_hook_rename,
|
||||
}
|
||||
|
||||
self._ida_hooks = CapaExplorerIdaHooks(self._ida_hook_screen_ea_changed, action_hooks)
|
||||
self._ida_hooks.hook()
|
||||
self.ida_hooks = CapaExplorerIdaHooks(self.ida_hook_screen_ea_changed, action_hooks)
|
||||
self.ida_hooks.hook()
|
||||
|
||||
def _unload_ida_hooks(self):
|
||||
''' unhook IDA user interface '''
|
||||
if self._ida_hooks:
|
||||
self._ida_hooks.unhook()
|
||||
def unload_ida_hooks(self):
|
||||
""" unhook IDA user interface """
|
||||
if self.ida_hooks:
|
||||
self.ida_hooks.unhook()
|
||||
|
||||
def _ida_hook_rename(self, meta, post=False):
|
||||
''' hook for IDA rename action
|
||||
def ida_hook_rename(self, meta, post=False):
|
||||
""" hook for IDA rename action
|
||||
|
||||
called twice, once before action and once after
|
||||
action completes
|
||||
|
||||
@param meta: TODO
|
||||
@param post: TODO
|
||||
'''
|
||||
ea = idaapi.get_screen_ea()
|
||||
if not ea or not capa.ida.helpers.is_func_start(ea):
|
||||
"""
|
||||
location = idaapi.get_screen_ea()
|
||||
if not location or not capa.ida.helpers.is_func_start(location):
|
||||
return
|
||||
|
||||
curr_name = idaapi.get_name(ea)
|
||||
curr_name = idaapi.get_name(location)
|
||||
|
||||
if post:
|
||||
# post action update data model w/ current name
|
||||
self._model_data.update_function_name(meta.get('prev_name', ''), curr_name)
|
||||
self.model_data.update_function_name(meta.get('prev_name', ''), curr_name)
|
||||
else:
|
||||
# pre action so save current name for replacement later
|
||||
meta['prev_name'] = curr_name
|
||||
|
||||
def _ida_hook_screen_ea_changed(self, widget, new_ea, old_ea):
|
||||
''' '''
|
||||
if not self._view_checkbox_limit_by.isChecked():
|
||||
def ida_hook_screen_ea_changed(self, widget, new_ea, old_ea):
|
||||
""" """
|
||||
if not self.view_checkbox_limit_by.isChecked():
|
||||
# ignore if checkbox not selected
|
||||
return
|
||||
|
||||
@@ -311,10 +331,10 @@ class CapaExplorerForm(idaapi.PluginForm):
|
||||
match = ''
|
||||
|
||||
# filter on virtual address to avoid updating filter string if function name is changed
|
||||
self._model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match)
|
||||
self.model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match)
|
||||
|
||||
def _load_capa_results(self):
|
||||
''' '''
|
||||
def load_capa_results(self):
|
||||
""" """
|
||||
logger.info('-' * 80)
|
||||
logger.info(' Using default embedded rules.')
|
||||
logger.info(' ')
|
||||
@@ -322,7 +342,7 @@ class CapaExplorerForm(idaapi.PluginForm):
|
||||
logger.info(' https://github.com/fireeye/capa-rules')
|
||||
logger.info('-' * 80)
|
||||
|
||||
rules_path = os.path.join(os.path.dirname(self._file_loc), '../..', 'rules')
|
||||
rules_path = os.path.join(os.path.dirname(self.file_loc), '../..', 'rules')
|
||||
rules = capa.main.get_rules(rules_path)
|
||||
rules = capa.rules.RuleSet(rules)
|
||||
capabilities = capa.main.find_capabilities(rules, capa.features.extractors.ida.IdaFeatureExtractor(), True)
|
||||
@@ -343,122 +363,154 @@ class CapaExplorerForm(idaapi.PluginForm):
|
||||
|
||||
capa.ida.helpers.inform_user_ida_ui('capa encountered warnings during analysis')
|
||||
|
||||
if capa.main.is_file_limitation(rules, capabilities, is_standalone=False):
|
||||
if capa.main.has_file_limitation(rules, capabilities, is_standalone=False):
|
||||
capa.ida.helpers.inform_user_ida_ui('capa encountered warnings during analysis')
|
||||
|
||||
logger.info('analysis completed.')
|
||||
|
||||
self._model_data.render_capa_results(rules, capabilities)
|
||||
self._render_capa_summary(rules, capabilities)
|
||||
doc = capa.render.convert_capabilities_to_result_document(rules, capabilities)
|
||||
|
||||
self._view_tree.sortByColumn(CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION, Qt.AscendingOrder)
|
||||
import json
|
||||
with open("C:\\Users\\spring\\Desktop\\hmm.json", "w") as twitter_data_file:
|
||||
json.dump(doc, twitter_data_file, indent=4, sort_keys=True, cls=capa.render.CapaJsonObjectEncoder)
|
||||
|
||||
self.model_data.render_capa_doc(doc)
|
||||
self.render_capa_doc_summary(doc)
|
||||
self.render_capa_doc_mitre_summary(doc)
|
||||
|
||||
self.view_tree.sortByColumn(CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION, QtCore.Qt.AscendingOrder)
|
||||
|
||||
logger.info('render views completed.')
|
||||
|
||||
def _render_capa_summary(self, ruleset, results):
|
||||
''' render results summary table
|
||||
def render_capa_doc_summary(self, doc):
|
||||
""" """
|
||||
for (row, rule) in enumerate(rutils.capability_rules(doc)):
|
||||
count = len(rule['matches'])
|
||||
|
||||
keep sync with capa.main
|
||||
if count == 1:
|
||||
capability = rule['meta']['name']
|
||||
else:
|
||||
capability = '%s (%d matches)' % (rule['meta']['name'], count)
|
||||
|
||||
@param ruleset: TODO
|
||||
@param results: TODO
|
||||
'''
|
||||
rules = set(filter(lambda x: not ruleset.rules[x].meta.get('lib', False), results.keys()))
|
||||
objectives = set()
|
||||
behaviors = set()
|
||||
techniques = set()
|
||||
self.view_summary.setRowCount(row + 1)
|
||||
|
||||
for rule in rules:
|
||||
parts = ruleset.rules[rule].meta.get(capa.main.RULE_CATEGORY, '').split('/')
|
||||
if len(parts) == 0 or list(parts) == ['']:
|
||||
continue
|
||||
if len(parts) > 0:
|
||||
objective = parts[0].replace('-', ' ')
|
||||
objectives.add(objective)
|
||||
if len(parts) > 1:
|
||||
behavior = parts[1].replace('-', ' ')
|
||||
behaviors.add(behavior)
|
||||
if len(parts) > 2:
|
||||
technique = parts[2].replace('-', ' ')
|
||||
techniques.add(technique)
|
||||
if len(parts) > 3:
|
||||
raise capa.rules.InvalidRule(capa.main.RULE_CATEGORY + ' tag must have at most three components')
|
||||
|
||||
# set row count to max set size
|
||||
self._view_summary.setRowCount(max(map(len, (rules, objectives, behaviors, techniques))))
|
||||
|
||||
# format rule hits
|
||||
rules = map(lambda x: '%s (%d)' % (x, len(results[x])), rules)
|
||||
|
||||
# sort results
|
||||
columns = list(map(lambda x: sorted(x, key=lambda s: s.lower()), (objectives, behaviors, techniques, rules)))
|
||||
|
||||
# load results into table by column
|
||||
for idx, column in enumerate(columns):
|
||||
self._load_view_summary_column(idx, column)
|
||||
self.view_summary.setItem(row, 0, self.render_new_table_header_item(capability))
|
||||
self.view_summary.setItem(row, 1, QtWidgets.QTableWidgetItem(rule['meta']['namespace']))
|
||||
|
||||
# resize columns to content
|
||||
self._view_summary.resizeColumnsToContents()
|
||||
self.view_summary.resizeColumnsToContents()
|
||||
|
||||
def _load_view_summary_column(self, column, texts):
|
||||
''' '''
|
||||
for row, text in enumerate(texts):
|
||||
self._view_summary.setItem(row, column, QTableWidgetItem(text))
|
||||
def render_capa_doc_mitre_summary(self, doc):
|
||||
""" """
|
||||
tactics = collections.defaultdict(set)
|
||||
for rule in rutils.capability_rules(doc):
|
||||
if not rule['meta'].get('att&ck'):
|
||||
continue
|
||||
|
||||
def _ida_reset(self):
|
||||
''' reset IDA user interface '''
|
||||
self._model_data.reset()
|
||||
self._view_tree.reset()
|
||||
self._view_checkbox_limit_by.setChecked(False)
|
||||
for attack in rule['meta']['att&ck']:
|
||||
tactic, _, rest = attack.partition('::')
|
||||
if '::' in rest:
|
||||
technique, _, rest = rest.partition('::')
|
||||
subtechnique, _, id = rest.rpartition(' ')
|
||||
tactics[tactic].add((technique, subtechnique, id))
|
||||
else:
|
||||
technique, _, id = rest.rpartition(' ')
|
||||
tactics[tactic].add((technique, id))
|
||||
|
||||
column_one = []
|
||||
column_two = []
|
||||
|
||||
for tactic, techniques in sorted(tactics.items()):
|
||||
column_one.append(tactic.upper())
|
||||
column_one.extend(['' for i in range(len(techniques) - 1)])
|
||||
|
||||
for spec in sorted(techniques):
|
||||
if len(spec) == 2:
|
||||
technique, id = spec
|
||||
column_two.append('%s %s' % (technique, id))
|
||||
elif len(spec) == 3:
|
||||
technique, subtechnique, id = spec
|
||||
column_two.append('%s::%s %s' % (technique, subtechnique, id))
|
||||
else:
|
||||
raise RuntimeError('unexpected ATT&CK spec format')
|
||||
|
||||
self.view_attack.setRowCount(max(len(column_one), len(column_two)))
|
||||
|
||||
for row, value in enumerate(column_one):
|
||||
self.view_attack.setItem(row, 0, self.render_new_table_header_item(value))
|
||||
|
||||
for row, value in enumerate(column_two):
|
||||
self.view_attack.setItem(row, 1, QtWidgets.QTableWidgetItem(value))
|
||||
|
||||
# resize columns to content
|
||||
self.view_attack.resizeColumnsToContents()
|
||||
|
||||
def render_new_table_header_item(self, text):
|
||||
""" """
|
||||
item = QtWidgets.QTableWidgetItem(text)
|
||||
item.setForeground(QtGui.QColor(88, 139, 174))
|
||||
|
||||
font = QtGui.QFont()
|
||||
font.setBold(True)
|
||||
|
||||
item.setFont(font)
|
||||
|
||||
return item
|
||||
|
||||
def ida_reset(self):
|
||||
""" reset IDA user interface """
|
||||
self.model_data.reset()
|
||||
self.view_tree.reset()
|
||||
self.view_checkbox_limit_by.setChecked(False)
|
||||
|
||||
def reload(self):
|
||||
''' reload views and re-run capa analysis '''
|
||||
self._ida_reset()
|
||||
self._model_proxy.invalidate()
|
||||
self._model_data.clear()
|
||||
self._view_summary.setRowCount(0)
|
||||
self._load_capa_results()
|
||||
""" reload views and re-run capa analysis """
|
||||
self.ida_reset()
|
||||
self.model_proxy.invalidate()
|
||||
self.model_data.clear()
|
||||
self.view_summary.setRowCount(0)
|
||||
self.load_capa_results()
|
||||
|
||||
logger.info('reload complete.')
|
||||
idaapi.info('%s reload completed.' % PLUGIN_NAME)
|
||||
|
||||
def reset(self):
|
||||
''' reset user interface elements
|
||||
""" reset user interface elements
|
||||
|
||||
e.g. checkboxes and IDA highlighting
|
||||
'''
|
||||
self._ida_reset()
|
||||
"""
|
||||
self.ida_reset()
|
||||
|
||||
logger.info('reset completed.')
|
||||
idaapi.info('%s reset completed.' % PLUGIN_NAME)
|
||||
|
||||
def _slot_menu_bar_hovered(self, action):
|
||||
''' display menu action tooltip
|
||||
def slot_menu_bar_hovered(self, action):
|
||||
""" display menu action tooltip
|
||||
|
||||
@param action: QAction*
|
||||
@param action: QtWidgets.QAction*
|
||||
|
||||
@reference: https://stackoverflow.com/questions/21725119/why-wont-qtooltips-appear-on-qactions-within-a-qmenu
|
||||
'''
|
||||
QToolTip.showText(QCursor.pos(), action.toolTip(), self._view_menu_bar, self._view_menu_bar.actionGeometry(action))
|
||||
"""
|
||||
QtWidgets.QToolTip.showText(QtGui.QCursor.pos(), action.toolTip(), self.view_menu_bar, self.view_menu_bar.actionGeometry(action))
|
||||
|
||||
def _slot_checkbox_limit_by_changed(self):
|
||||
''' slot activated if checkbox clicked
|
||||
def slot_checkbox_limit_by_changed(self):
|
||||
""" slot activated if checkbox clicked
|
||||
|
||||
if checked, configure function filter if screen ea is located
|
||||
if checked, configure function filter if screen location is located
|
||||
in function, otherwise clear filter
|
||||
'''
|
||||
"""
|
||||
match = ''
|
||||
if self._view_checkbox_limit_by.isChecked():
|
||||
ea = capa.ida.helpers.get_func_start_ea(idaapi.get_screen_ea())
|
||||
if ea:
|
||||
match = capa.ida.explorer.item.ea_to_hex_str(ea)
|
||||
self._model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match)
|
||||
if self.view_checkbox_limit_by.isChecked():
|
||||
location = capa.ida.helpers.get_func_start_ea(idaapi.get_screen_ea())
|
||||
if location:
|
||||
match = capa.ida.explorer.item.location_to_hex(location)
|
||||
self.model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match)
|
||||
|
||||
self._view_tree.resize_columns_to_content()
|
||||
self.view_tree.resize_columns_to_content()
|
||||
|
||||
|
||||
def main():
|
||||
''' TODO: move to idaapi.plugin_t class '''
|
||||
""" TODO: move to idaapi.plugin_t class """
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
if not capa.ida.helpers.is_supported_file_type():
|
||||
|
||||
Reference in New Issue
Block a user