diff --git a/capa/ida/explorer/item.py b/capa/ida/explorer/item.py index d49a5c3f..729faa4a 100644 --- a/capa/ida/explorer/item.py +++ b/capa/ida/explorer/item.py @@ -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) diff --git a/capa/ida/explorer/model.py b/capa/ida/explorer/model.py index 06b64e86..7b3b21c7 100644 --- a/capa/ida/explorer/model.py +++ b/capa/ida/explorer/model.py @@ -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="")" 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) diff --git a/capa/ida/explorer/proxy.py b/capa/ida/explorer/proxy.py index 465abd5a..811eff16 100644 --- a/capa/ida/explorer/proxy.py +++ b/capa/ida/explorer/proxy.py @@ -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) diff --git a/capa/ida/explorer/view.py b/capa/ida/explorer/view.py index 2868d01a..d2a8eb46 100644 --- a/capa/ida/explorer/view.py +++ b/capa/ida/explorer/view.py @@ -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) diff --git a/capa/ida/ida_capa_explorer.py b/capa/ida/ida_capa_explorer.py index fb2759ec..873e0513 100644 --- a/capa/ida/ida_capa_explorer.py +++ b/capa/ida/ida_capa_explorer.py @@ -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():