def __init__(self, parent=None):
""" """
super(CapaExplorerDataModel, self).__init__(parent)
self.root_node = CapaExplorerDataItem(None, ["Rule Information", "Address", "Details"])
class CapaExplorerDataModel(QtCore.QAbstractItemModel):
""" """
COLUMN_INDEX_RULE_INFORMATION = 0
COLUMN_INDEX_VIRTUAL_ADDRESS = 1
COLUMN_INDEX_DETAILS = 2
COLUMN_COUNT = 3
def __init__(self, parent=None):
""" """
super(CapaExplorerDataModel, self).__init__(parent)
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_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.reset_ida_highlighting(model_index.internalPointer(), False)
self.dataChanged.emit(model_index, model_index)
def clear(self):
""" """
self.beginResetModel()
self.root_node.removeChildren()
self.endResetModel()
def columnCount(self, model_index):
"""get the number of columns for the children of the given parent
@param model_index: QModelIndex*
@retval column count
"""
if model_index.isValid():
return model_index.internalPointer().columnCount()
else:
return self.root_node.columnCount()
def data(self, model_index, role):
"""get data stored under the given role for the item referred to by the index
@param model_index: QModelIndex*
@param role: QtCore.Qt.*
@retval data to be displayed
"""
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 item.data(column)
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 column == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION:
# inform view how to display content of checkbox - un/checked
return QtCore.Qt.Checked if item.isChecked() else QtCore.Qt.Unchecked
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 isinstance(
item,
(
CapaExplorerRuleItem,
CapaExplorerRuleMatchItem,
CapaExplorerBlockItem,
CapaExplorerFunctionItem,
CapaExplorerFeatureItem,
CapaExplorerSubscopeItem,
),
)
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, model_index):
"""get item flags for given index
@param model_index: QModelIndex*
@retval QtCore.Qt.ItemFlags
"""
if not model_index.isValid():
return QtCore.Qt.NoItemFlags
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
@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_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
@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_node
else:
parent_item = parent.internalPointer()
child_item = parent_item.child(row)
if child_item:
return self.createIndex(row, column, child_item)
else:
return QtCore.QModelIndex()
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 model_index: QModelIndex*
@retval QModelIndex*
"""
if not model_index.isValid():
return QtCore.QModelIndex()
child = model_index.internalPointer()
parent = child.parent()
if parent == self.root_node:
return QtCore.QModelIndex()
return self.createIndex(parent.row(), 0, parent)
def iterateChildrenIndexFromRootIndex(self, model_index, ignore_root=True):
"""depth-first traversal of child nodes
@param model_index: QModelIndex*
@param ignore_root: if set, do not return root index
@retval yield QModelIndex*
"""
visited = set()
stack = deque((model_index,))
while True:
try:
child_index = stack.pop()
except IndexError:
break
if child_index not in visited:
if not ignore_root or child_index is not model_index:
# ignore root
yield child_index
visited.add(child_index)
for idx in range(self.rowCount(child_index)):
stack.append(child_index.child(idx, 0))
def reset_ida_highlighting(self, item, checked):
"""reset IDA highlight for an item
@param item: capa explorer item
@param checked: indicates item is or not checked
"""
if not isinstance(
item, (CapaExplorerStringViewItem, CapaExplorerInstructionViewItem, CapaExplorerByteViewItem)
):
# ignore other item types
return
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.location, idc.CIC_ITEM, DEFAULT_HIGHLIGHT)
else:
# item unchecked - reset highlight
if curr_highlight != DEFAULT_HIGHLIGHT:
# user modified highlight - record new highlight and do not modify
item.ida_highlight = curr_highlight
else:
# reset highlight to previous
idc.set_color(item.location, idc.CIC_ITEM, item.ida_highlight)
def setData(self, model_index, value, role):
"""set the role data for the item at index to value
@param model_index: QModelIndex*
@param value: QVariant*
@param role: QtCore.Qt.EditRole
@retval True/False
"""
if not model_index.isValid():
return False
if (
role == QtCore.Qt.CheckStateRole
and model_index.column() == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION
):
# user un/checked box - un/check parent and children
for child_index in self.iterateChildrenIndexFromRootIndex(model_index, ignore_root=False):
child_index.internalPointer().setChecked(value)
self.reset_ida_highlighting(child_index.internalPointer(), value)
self.dataChanged.emit(child_index, child_index)
return True
if (
role == QtCore.Qt.EditRole
and value
and 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 = model_index.internalPointer().info
new_name = str(value)
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
# no handle
return False
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 model_index: QModelIndex*
@retval row count
"""
if model_index.column() > 0:
return 0
if not model_index.isValid():
item = self.root_node
else:
item = model_index.internalPointer()
return item.childCount()
def render_capa_doc_statement_node(self, parent, statement, locations, doc):
"""render capa statement read from doc
@param parent: parent to which new child is assigned
@param statement: statement read from doc
@param locations: locations of children (applies to range only?)
@param doc: capa result doc
"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, str(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"])
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"])
parent2 = CapaExplorerFeatureItem(parent, display=display)
for location in locations:
# for each location render child node for range statement
self.render_capa_doc_feature(parent2, statement["child"], location, doc)
return parent2
elif statement["type"] == "subscope":
return CapaExplorerSubscopeItem(parent, statement[statement["type"]])
else:
raise RuntimeError("unexpected match statement type: " + str(statement))
def render_capa_doc_match(self, parent, match, doc):
"""render capa match read from doc
@param parent: parent node to which new child is assigned
@param match: match read from doc
@param doc: capa result doc
"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
# 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"], match.get("locations", []), doc
)
elif match["node"]["type"] == "feature":
parent2 = self.render_capa_doc_feature_node(
parent, match["node"]["feature"], match.get("locations", []), doc
)
else:
raise RuntimeError("unexpected node type: " + str(match["node"]["type"]))
for child in match.get("children", []):
self.render_capa_doc_match(parent2, child, doc)
def render_capa_doc(self, doc):
"""render capa features specified in doc
@param doc: capa result doc
"""
# inform model that changes are about to occur
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["rules"][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)
# inform model changes have ended
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
Example:
"feature": {
"bytes": "01 14 02 00 00 00 00 00 C0 00 00 00 00 00 00 46",
"description": "CLSID_ShellLink",
"type": "bytes"
}
bytes(01 14 02 00 00 00 00 00 C0 00 00 00 00 00 00 46 = CLSID_ShellLink)
"""
if feature[feature["type"]]:
if feature.get("description", ""):
return "%s(%s = %s)" % (feature["type"], feature[feature["type"]], feature["description"])
else:
return "%s(%s)" % (feature["type"], feature[feature["type"]])
else:
return "%s" % feature["type"]
def render_capa_doc_feature_node(self, parent, feature, locations, doc):
"""process capa doc feature node
@param parent: parent node to which child is assigned
@param feature: capa doc feature node
@param locations: locations identified for feature
@param doc: capa doc
Example:
"feature": {
"description": "FILE_WRITE_DATA",
"number": "0x2",
"type": "number"
}
"""
display = self.capa_doc_feature_to_display(feature)
if len(locations) == 1:
# only one location for feature so no need to nest children
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
@param location: address of feature
@param display: text to display in plugin ui
Example:
"feature": {
"description": "FILE_WRITE_DATA",
"number": "0x2",
"type": "number"
}
"""
# special handling for characteristic pending type
if feature["type"] == "characteristic":
if feature[feature["type"]] in ("embedded pe",):
return CapaExplorerByteViewItem(parent, display, location)
if feature[feature["type"]] in ("loop", "recursive call", "tight loop"):
return CapaExplorerFeatureItem(parent, display=display)
# default to instruction view for all other characteristics
return CapaExplorerInstructionViewItem(parent, display, location)
if feature["type"] == "match":
# display content of rule for all rule matches
return CapaExplorerRuleMatchItem(
parent, display, source=doc["rules"].get(feature[feature["type"]], {}).get("source", "")
)
if feature["type"] == "regex":
return CapaExplorerFeatureItem(parent, display, location, details=feature["match"])
if feature["type"] == "basicblock":
return CapaExplorerBlockItem(parent, location)
if feature["type"] in (
"bytes",
"api",
"mnemonic",
"number",
"offset",
"number/x32",
"number/x64",
"offset/x32",
"offset/x64",
):
# display instruction preview
return CapaExplorerInstructionViewItem(parent, display, location)
if feature["type"] in ("section",):
# display byte preview
return CapaExplorerByteViewItem(parent, display, location)
if feature["type"] in ("string",):
# display string preview
return CapaExplorerStringViewItem(parent, display, location)
if feature["type"] in ("import", "export"):
# display no preview
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 old function name with new function name
@param old_name: previous function name
@param new_name: new function name
"""
# create empty root index for search
root_index = self.index(0, 0, QtCore.QModelIndex())
# convert name to view format for matching e.g. function(my_function)
old_name = CapaExplorerFunctionItem.fmt % old_name
# recursive search for all instances of old function name
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
# replace old function name with new function name and emit change
model_index.internalPointer().info = new_name
self.dataChanged.emit(model_index, model_index)