gitea-mirror/capa
1
0
Fork 0
mirror of https://github.com/mandiant/capa.git synced 2026-02-04 11:07:53 -08:00
Code Issues Packages Projects Releases Wiki Activity

rulegen adding support for sync between editor and preview windows

Browse Source
This commit is contained in:
Michael Hunhoff
2021-01-28 17:15:18 -07:00
parent 9caafedb8d
commit b413f2eafe
2 changed files with 312 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
capa/ida/plugin/form.py
View File

@@ -7,15 +7,12 @@
# See the License for the specific language governing permissions and limitations under the License.
import os
import sys
import json
import types
import copy
import logging
import itertools
import collections
import idc
import idaapi
import ida_kernwin
import ida_settings
@@ -57,13 +54,6 @@ def trim_function_name(f, max_length=25):
return n
def trim_scope(n):
""" """
if "/" in n:
n = n.rpartition("/")[0]
return n
def find_func_features(f, extractor):
""" """
func_features = collections.defaultdict(set)
@@ -392,6 +382,7 @@ class CapaExplorerForm(idaapi.PluginForm):
self.view_rulegen_features = CapaExplorerRulegenFeatures(self.view_rulegen_editor, parent=self.parent)
self.view_rulegen_preview.textChanged.connect(self.slot_rulegen_preview_update)
self.view_rulegen_editor.updated.connect(self.slot_rulegen_editor_update)
self.set_rulegen_preview_border_neutral()
@@ -420,7 +411,7 @@ class CapaExplorerForm(idaapi.PluginForm):
("Change default rule author...", "Set default rule author", self.slot_change_rule_author),
("Change default rule scope...", "Set default rule scope", self.slot_change_rule_scope),
)
self.load_menu("Configuration", actions)
self.load_menu("Settings", actions)
def load_menu(self, title, actions):
"""load menu actions
@@ -858,9 +849,9 @@ class CapaExplorerForm(idaapi.PluginForm):
""" """
self.view_rulegen_preview.setStyleSheet("border: 3px solid green")
def slot_rulegen_preview_update(self):
def update_rule_status(self, rule_text):
""" """
if not self.view_rulegen_editor.root:
if self.view_rulegen_editor.root is None:
self.set_rulegen_preview_border_neutral()
self.view_rulegen_status_label.clear()
return
@@ -868,7 +859,7 @@ class CapaExplorerForm(idaapi.PluginForm):
self.set_rulegen_preview_border_error()
try:
rule = capa.rules.Rule.from_yaml(self.view_rulegen_preview.toPlainText())
rule = capa.rules.Rule.from_yaml(rule_text)
except Exception as e:
self.set_rulegen_status("Failed to compile rule! %s" % e)
return
@@ -901,6 +892,17 @@ class CapaExplorerForm(idaapi.PluginForm):
"Rule compiled, but no match found for %s" % idaapi.get_name(self.rulegen_current_function.start_ea)
)
def slot_rulegen_editor_update(self):
""" """
rule_text = self.view_rulegen_preview.toPlainText()
self.update_rule_status(rule_text)
def slot_rulegen_preview_update(self):
""" """
rule_text = self.view_rulegen_preview.toPlainText()
self.view_rulegen_editor.load_features_from_yaml(rule_text, False)
self.update_rule_status(rule_text)
def slot_limit_rulegen_features_to_search(self, text):
""" """
self.view_rulegen_features.filter_items_by_text(text)
@@ -952,8 +954,8 @@ class CapaExplorerForm(idaapi.PluginForm):
if not s:
idaapi.info("No rule to save.")
return
path = idaapi.ask_file(True, "*.yml", "Choose file to save capa rule")
path = self.ask_user_capa_rule_file()
if not path:
return
@@ -1003,6 +1005,12 @@ class CapaExplorerForm(idaapi.PluginForm):
)
)
def ask_user_capa_rule_file(self):
""" """
return QtWidgets.QFileDialog.getSaveFileName(
None, "Please select a capa rule to edit", settings.user["rule_path"], "*.yml"
)[0]
def slot_change_rule_scope(self):
""" """
scope = idaapi.ask_str(str(settings.user.get("rulegen_scope", "function")), 0, "Enter default rule scope")

413
capa/ida/plugin/view.py
View File

@@ -5,8 +5,7 @@
# Unless required by applicable law or agreed to in writing, software distributed under the License
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.
from collections import Counter, defaultdict
import binascii
from collections import Counter
import re
import idc
@@ -26,6 +25,116 @@ COLOR_GREEN_RGB = (79, 121, 66)
COLOR_BLUE_RGB = (37, 147, 215)
def calc_level_by_indent(line, prev_level=0):
""" """
if not len(line.strip()):
# blank line, which may occur for comments so we simply use the last level
return prev_level
stripped = line.lstrip()
if stripped.startswith("description"):
# need to adjust two spaces when encountering string description
line = line[2:]
# calc line level based on preceding whitespace
return len(line) - len(stripped)
def parse_feature_for_node(feature):
""" """
description = ""
comment = ""
if feature.startswith("- count"):
# count is weird, we need to handle special
# first, we need to grab the comment, if exists
# next, we need to check for an embedded description
feature, _, comment = feature.partition("#")
m = re.search(r"- count\(([a-zA-Z]+)\((.+)\s+=\s+(.+)\)\):\s*(.+)", feature)
if m:
# reconstruct count without description
feature, value, description, count = m.groups()
feature = "- count(%s(%s)): %s" % (feature, value, count)
elif not feature.startswith("#"):
feature, _, comment = feature.partition("#")
feature, _, description = feature.partition("=")
return map(str.strip, (feature, description, comment))
def parse_node_for_feature(feature, description, comment, depth):
""" """
depth = (depth * 2) + 4
display = ""
if feature.startswith("#"):
display += "%s%s\n" % (" " * depth, feature)
elif description:
if feature.startswith(("- and", "- or", "- optional", "- basic block", "- not")):
display += "%s%s" % (" " * depth, feature)
if comment:
display += " # %s" % comment
display += "\n%s- description: %s\n" % (" " * (depth + 2), description)
elif feature.startswith("- string"):
display += "%s%s" % (" " * depth, feature)
if comment:
display += " # %s" % comment
display += "\n%sdescription: %s\n" % (" " * (depth + 2), description)
elif feature.startswith("- count"):
# count is weird, we need to format description based on feature type, so we parse with regex
# assume format - count(<feature_name>(<feature_value>)): <count>
m = re.search(r"- count\(([a-zA-Z]+)\((.+)\)\): (.+)", feature)
if m:
name, value, count = m.groups()
if name in ("string",):
display += "%s%s" % (" " * depth, feature)
if comment:
display += " # %s" % comment
display += "\n%sdescription: %s\n" % (" " * (depth + 2), description)
else:
display += "%s- count(%s(%s = %s)): %s" % (
" " * depth,
name,
value,
description,
count,
)
if comment:
display += " # %s\n" % comment
else:
display += "%s%s = %s" % (" " * depth, feature, description)
if comment:
display += " # %s\n" % comment
else:
display += "%s%s" % (" " * depth, feature)
if comment:
display += " # %s\n" % comment
return display if display.endswith("\n") else display + "\n"
def yaml_to_nodes(s):
level = 0
for line in s.splitlines():
feature, description, comment = parse_feature_for_node(line.strip())
o = QtWidgets.QTreeWidgetItem(None)
# set node attributes
setattr(o, "capa_level", calc_level_by_indent(line, level))
if feature.startswith(("- and:", "- or:", "- not:", "- basic block:", "- optional:")):
setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_expression())
elif feature.startswith("#"):
setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_comment())
else:
setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_feature())
# set node text
for (i, v) in enumerate((feature, description, comment)):
o.setText(i, v)
yield o
def iterate_tree(o):
""" """
itr = QtWidgets.QTreeWidgetItemIterator(o)
@@ -71,6 +180,8 @@ class CapaExplorerRulgenPreview(QtWidgets.QTextEdit):
super(CapaExplorerRulgenPreview, self).__init__(parent)
self.setFont(QtGui.QFont("Courier", weight=QtGui.QFont.Bold))
self.setLineWrapMode(QtWidgets.QTextEdit.NoWrap)
self.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarAsNeeded)
def reset_view(self):
""" """
@@ -93,16 +204,27 @@ class CapaExplorerRulgenPreview(QtWidgets.QTextEdit):
]
self.setText("\n".join(metadata_default))
def keyPressEvent(self, e):
""" """
if e.key() == QtCore.Qt.Key_Tab:
self.insertPlainText(" " * 2)
else:
super(CapaExplorerRulgenPreview, self).keyPressEvent(e)
class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
updated = QtCore.pyqtSignal()
def __init__(self, preview, parent=None):
""" """
super(CapaExplorerRulgenEditor, self).__init__(parent)
self.preview = preview
self.setHeaderLabels(["Feature", "Description"])
self.setHeaderLabels(["Feature", "Description", "Comment"])
self.header().setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents)
self.header().setStretchLastSection(False)
self.setExpandsOnDoubleClick(False)
self.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
@@ -134,6 +256,26 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
""" """
return 1
@staticmethod
def get_column_comment_index():
""" """
return 2
@staticmethod
def get_node_type_expression():
""" """
return 0
@staticmethod
def get_node_type_feature():
""" """
return 1
@staticmethod
def get_node_type_comment():
""" """
return 2
def dragMoveEvent(self, e):
""" """
super(CapaExplorerRulgenEditor, self).dragMoveEvent(e)
@@ -167,22 +309,17 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
def slot_remove_selected(self, action):
""" """
for o in self.selectedItems():
# do not remove root node from tree
if o == self.root:
self.takeTopLevelItem(self.indexOfTopLevelItem(o))
self.root = None
continue
o.parent().removeChild(o)
def slot_nest_features(self, action):
""" """
# create a new parent under root node, by default; new node added last position in tree
new_parent = self.add_child_item(
self.root,
(action.data()[0], ""),
drop_enabled=True,
select_enabled=True,
drag_enabled=True,
has_children=True,
)
new_parent = self.new_expression_node(self.root, (action.data()[0], ""))
for o in self.get_features(selected=True):
# take child from its parent by index, add to new parent
new_parent.addChild(o.parent().takeChild(o.parent().indexOfChild(o)))
@@ -204,73 +341,50 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
if not self.indexAt(pos).isValid():
# user selected invalid index
self.load_custom_context_menu_invalid_index(pos)
elif not self.itemAt(pos).flags() & QtCore.Qt.ItemIsEditable:
elif self.itemAt(pos).capa_type == CapaExplorerRulgenEditor.get_node_type_expression():
# user selected expression node
self.load_custom_context_menu_expression(pos)
else:
# user selected feature node
self.load_custom_context_menu_feature(pos)
# refresh views
# self.prune_expressions()
self.update_preview()
def slot_item_double_clicked(self, o, column):
""" """
if o.flags() & QtCore.Qt.ItemIsEditable:
if column in (
CapaExplorerRulgenEditor.get_column_comment_index(),
CapaExplorerRulgenEditor.get_column_description_index(),
):
o.setFlags(o.flags() | QtCore.Qt.ItemIsEditable)
self.editItem(o, column)
o.setFlags(o.flags() & ~QtCore.Qt.ItemIsEditable)
self.is_editing = True
else:
if column == CapaExplorerRulgenEditor.get_column_description_index():
o.setFlags(o.flags() | QtCore.Qt.ItemIsEditable)
self.editItem(o, column)
o.setFlags(o.flags() & ~QtCore.Qt.ItemIsEditable)
self.is_editing = True
def update_preview(self):
""" """
rule_text = self.preview.toPlainText()
rule_text = rule_text[: rule_text.find("features:") + len("features:")]
rule_text += "\n"
if -1 != rule_text.find("features:"):
rule_text = rule_text[: rule_text.find("features:") + len("features:")]
rule_text += "\n"
else:
rule_text = rule_text.rstrip()
rule_text += "\n features:\n"
for o in iterate_tree(self):
display = o.text(CapaExplorerRulgenEditor.get_column_feature_index())
description = o.text(CapaExplorerRulgenEditor.get_column_description_index())
depth_space = (calc_item_depth(o) * 2) + 4
if not description:
rule_text += "%s%s\n" % (" " * depth_space, display)
else:
if display.startswith(("- and", "- or", "- optional", "- basic block", "- not")):
rule_text += "%s%s\n" % (" " * depth_space, display)
rule_text += "%s- description: %s\n" % (" " * (depth_space + 2), description)
continue
elif display.startswith("- string"):
rule_text += "%s%s\n" % (" " * depth_space, display)
rule_text += "%sdescription: %s\n" % (" " * (depth_space + 2), description)
continue
elif display.startswith("- count"):
# count is weird, we need to format description based on feature type, so we parse with regex
# assume format - count(<feature_name>(<feature_value>)): <count>
m = re.search(r"- count\(([a-zA-Z]+)\((.+)\)\): (.+)", display)
if m:
name, value, count = m.groups()
if name in ("string",):
rule_text += "%s%s\n" % (" " * depth_space, display)
rule_text += "%sdescription: %s\n" % (" " * (depth_space + 2), description)
else:
rule_text += "%s- count(%s(%s = %s)): %s\n" % (
" " * depth_space,
name,
value,
description,
count,
)
continue
# catchall
rule_text += "%s%s = %s\n" % (" " * depth_space, display, description)
feature, description, comment = map(str.strip, tuple(o.text(i) for i in range(3)))
rule_text += parse_node_for_feature(feature, description, comment, calc_item_depth(o))
# FIXME we avoid circular update by disabling signals when updating
# the preview. Preferably we would refactor the code to avoid this
# in the first place
self.preview.blockSignals(True)
self.preview.setPlainText(rule_text)
self.preview.blockSignals(False)
# emit signal so views can update
self.updated.emit()
def load_custom_context_menu_invalid_index(self, pos):
""" """
@@ -281,8 +395,6 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
def load_custom_context_menu_feature(self, pos):
""" """
# sub_sub_actions = []
actions = (("Remove selection", (), self.slot_remove_selected),)
sub_actions = (
@@ -295,17 +407,8 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
feature_count = len(tuple(self.get_features(selected=True)))
"""
for i in range(feature_count + 1):
sub_sub_actions.append(("%d or more" % i, ("- %d or more:" % i,), self.slot_nest_features))
sub_sub_menu = build_custom_context_menu(self.parent(), sub_sub_actions)
sub_sub_menu.setTitle("N or more")
"""
sub_menu = build_custom_context_menu(self.parent(), sub_actions)
sub_menu.setTitle("Nest feature%s" % ("" if feature_count == 1 else "s"))
# sub_menu.addMenu(sub_sub_menu)
menu = build_custom_context_menu(self.parent(), actions)
menu.addMenu(sub_menu)
@@ -320,30 +423,27 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
("not", ("- not:", self.itemAt(pos)), self.slot_edit_expression),
("optional", ("- optional:", self.itemAt(pos)), self.slot_edit_expression),
("basic block", ("- basic block:", self.itemAt(pos)), self.slot_edit_expression),
("Remove expression", (), self.slot_remove_selected),
)
actions = []
actions = (("Remove expression", (), self.slot_remove_selected),)
sub_menu = build_custom_context_menu(self.parent(), sub_actions)
sub_menu.setTitle("Modify")
if self.root != self.itemAt(pos):
# only add remove option if not root
actions.append(("Remove expression", (), self.slot_remove_selected))
menu = build_custom_context_menu(self.parent(), actions)
menu.addMenu(sub_menu)
menu.exec_(self.viewport().mapToGlobal(pos))
def style_parent_node(self, o):
def style_expression_node(self, o):
""" """
font = QtGui.QFont()
font.setBold(True)
o.setFont(CapaExplorerRulgenEditor.get_column_feature_index(), font)
def style_child_node(self, o):
def style_feature_node(self, o):
""" """
font = QtGui.QFont()
brush = QtGui.QBrush()
@@ -355,77 +455,138 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
o.setFont(CapaExplorerRulgenEditor.get_column_feature_index(), font)
o.setForeground(CapaExplorerRulgenEditor.get_column_feature_index(), brush)
def add_child_item(
self,
parent,
values,
data=None,
drop_enabled=False,
edit_enabled=False,
select_enabled=False,
drag_enabled=False,
has_children=False,
):
def style_comment_node(self, o):
""" """
c = QtWidgets.QTreeWidgetItem(parent)
font = QtGui.QFont()
font.setBold(True)
font.setFamily("Courier")
if has_children:
# adding expression node, set bold
self.style_parent_node(c)
else:
# adding feature node, set style, weight, and color
self.style_child_node(c)
o.setFont(CapaExplorerRulgenEditor.get_column_feature_index(), font)
if not select_enabled:
c.setFlags(c.flags() & ~QtCore.Qt.ItemIsSelectable)
if edit_enabled:
c.setFlags(c.flags() | QtCore.Qt.ItemIsTristate | QtCore.Qt.ItemIsEditable)
if not drop_enabled:
c.setFlags(c.flags() & ~QtCore.Qt.ItemIsDropEnabled)
if drag_enabled:
c.setFlags(c.flags() | QtCore.Qt.ItemIsDragEnabled)
def set_expression_node(self, o):
""" """
setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_expression())
self.style_expression_node(o)
def set_feature_node(self, o):
""" """
setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_feature())
o.setFlags(o.flags() & ~QtCore.Qt.ItemIsDropEnabled)
self.style_feature_node(o)
def set_comment_node(self, o):
""" """
setattr(o, "capa_type", CapaExplorerRulgenEditor.get_node_type_comment())
o.setFlags(o.flags() & ~QtCore.Qt.ItemIsDropEnabled)
self.style_comment_node(o)
def new_expression_node(self, parent, values=()):
""" """
o = QtWidgets.QTreeWidgetItem(parent)
self.set_expression_node(o)
for (i, v) in enumerate(values):
c.setText(i, v)
o.setText(i, v)
return o
if data:
c.setData(0, 0x100, data)
def new_feature_node(self, parent, values=()):
""" """
o = QtWidgets.QTreeWidgetItem(parent)
self.set_feature_node(o)
for (i, v) in enumerate(values):
o.setText(i, v)
return o
return c
def new_comment_node(self, parent, values=()):
""" """
o = QtWidgets.QTreeWidgetItem(parent)
self.set_comment_node(o)
for (i, v) in enumerate(values):
o.setText(i, v)
return o
def update_features(self, features):
""" """
if not self.root:
# root node does not exist, create default node, set expanded
self.root = self.add_child_item(
self, ("- or:", ""), drop_enabled=True, select_enabled=True, has_children=True
)
self.root.setExpanded(True)
self.root = self.new_expression_node(self, ("- or:", ""))
# build feature counts
counted = list(zip(Counter(features).keys(), Counter(features).values()))
# single features
for (k, v) in filter(lambda t: t[1] == 1, counted):
display = "- %s: %s" % (k.name.lower(), k.get_value_str())
self.add_child_item(self.root, (display, ""), edit_enabled=True, select_enabled=True, drag_enabled=True)
self.new_feature_node(self.root, ("- %s: %s" % (k.name.lower(), k.get_value_str()), ""))
# n > 1 features
for (k, v) in filter(lambda t: t[1] > 1, counted):
display = "- count(%s): %d" % (str(k), v)
self.add_child_item(self.root, (display, ""), edit_enabled=True, select_enabled=True, drag_enabled=True)
self.new_feature_node(self.root, ("- count(%s): %d" % (str(k), v), ""))
self.expandAll()
self.update_preview()
def prune_expressions(self):
def load_features_from_yaml(self, rule_text, update_preview=False):
""" """
for o in self.get_expressions(ignore=(self.root,)):
if not o.childCount():
o.parent().removeChild(o)
def add_node(parent, node):
if node.text(0).startswith("description:"):
if parent.childCount():
parent.child(parent.childCount() - 1).setText(1, node.text(0).lstrip("description:").lstrip())
else:
parent.setText(1, node.text(0).lstrip("description:").lstrip())
elif node.text(0).startswith("- description:"):
parent.setText(1, node.text(0).lstrip("- description:").lstrip())
else:
parent.addChild(node)
def build(parent, nodes):
if nodes:
child_lvl = nodes[0].capa_level
while nodes:
node = nodes.pop(0)
if node.capa_level == child_lvl:
add_node(parent, node)
elif node.capa_level > child_lvl:
nodes.insert(0, node)
build(parent.child(parent.childCount() - 1), nodes)
else:
parent = parent.parent() if parent.parent() else parent
add_node(parent, node)
self.reset_view()
# check for lack of features block
if -1 == rule_text.find("features:"):
return
rule_features = rule_text[rule_text.find("features:") + len("features:") :].strip()
rule_nodes = list(yaml_to_nodes(rule_features))
# check for lack of nodes
if not rule_nodes:
return
for o in rule_nodes:
(self.set_expression_node, self.set_feature_node, self.set_comment_node)[o.capa_type](o)
self.root = rule_nodes.pop(0)
self.addTopLevelItem(self.root)
if update_preview:
self.preview.blockSignals(True)
self.preview.setPlainText(rule_text)
self.preview.blockSignals(False)
build(self.root, rule_nodes)
self.expandAll()
def get_features(self, selected=False, ignore=()):
""" """
for feature in filter(lambda o: o.flags() & QtCore.Qt.ItemIsEditable, tuple(iterate_tree(self))):
for feature in filter(
lambda o: o.capa_type
in (CapaExplorerRulgenEditor.get_node_type_feature(), CapaExplorerRulgenEditor.get_node_type_comment()),
tuple(iterate_tree(self)),
):
if feature in ignore:
continue
if selected and not feature.isSelected():
@@ -434,7 +595,9 @@ class CapaExplorerRulgenEditor(QtWidgets.QTreeWidget):
def get_expressions(self, selected=False, ignore=()):
""" """
for expression in filter(lambda o: not o.flags() & QtCore.Qt.ItemIsEditable, tuple(iterate_tree(self))):
for expression in filter(
lambda o: o.capa_type == CapaExplorerRulgenEditor.get_node_type_expression(), tuple(iterate_tree(self))
):
if expression in ignore:
continue
if selected and not expression.isSelected():
@@ -527,14 +690,14 @@ class CapaExplorerRulegenFeatures(QtWidgets.QTreeWidget):
o.setHidden(False)
o.setExpanded(True)
def style_parent_node(self, o):
def set_expression_node(self, o):
""" """
font = QtGui.QFont()
font.setBold(True)
o.setFont(CapaExplorerRulegenFeatures.get_column_feature_index(), font)
def style_child_node(self, o):
def style_feature_node(self, o):
""" """
font = QtGui.QFont("Courier", weight=QtGui.QFont.Bold)
brush = QtGui.QBrush()
@@ -553,10 +716,10 @@ class CapaExplorerRulegenFeatures(QtWidgets.QTreeWidget):
c = QtWidgets.QTreeWidgetItem(parent)
if has_children:
self.style_parent_node(c)
self.set_expression_node(c)
c.setFlags(c.flags() & ~QtCore.Qt.ItemIsSelectable)
else:
self.style_child_node(c)
self.style_feature_node(c)
for (i, v) in enumerate(values):
c.setText(i, v)