Merge pull request #58 from fireeye/capa-explorer-support-doc-format

Capa explorer support doc format
This commit is contained in:
Willi Ballenthin
2020-07-01 09:50:42 -06:00
committed by GitHub
5 changed files with 846 additions and 656 deletions

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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():