From 9c3e6cd3cd6afb4ddc95cf992e46d58efc29e2d4 Mon Sep 17 00:00:00 2001 From: William Ballenthin Date: Tue, 14 Jul 2020 10:14:19 -0600 Subject: [PATCH] remove IDA rule generator closes #134 --- capa/ida/ida_rule_generator.py | 288 --------------------------------- 1 file changed, 288 deletions(-) delete mode 100644 capa/ida/ida_rule_generator.py diff --git a/capa/ida/ida_rule_generator.py b/capa/ida/ida_rule_generator.py deleted file mode 100644 index 1f26d325..00000000 --- a/capa/ida/ida_rule_generator.py +++ /dev/null @@ -1,288 +0,0 @@ -# TODO documentation - -import logging -import binascii -import textwrap -from collections import Counter, defaultdict - -import idc -import idaapi -from PyQt5 import QtCore, QtWidgets -from PyQt5.QtWidgets import QTextEdit, QHeaderView, QTreeWidget, QTreeWidgetItem - -import capa -import capa.main -import capa.features.extractors.ida.helpers -from capa.ida import plugin_helpers - -logger = logging.getLogger("rulegen") - - -AUTHOR_NAME = "" -COLOR_HIGHLIGHT = 0xD096FF - - -def get_func_start(ea): - f = idaapi.get_func(ea) - if f: - return f.start_ea - else: - return None - - -class Hooks(idaapi.UI_Hooks): - """ - Notifies the plugin when navigating to another function - NOTE: it uses the global variable FLEX to access the - PluginForm object. This looks nasty, maybe there is a better way? - """ - - def screen_ea_changed(self, ea, prev_ea): - widget = idaapi.get_current_widget() - if idaapi.get_widget_type(widget) != idaapi.BWN_DISASM: - # Ignore non disassembly views - return - - try: - f1 = get_func_start(ea) - f2 = get_func_start(prev_ea) - - if f1 != f2: - # changed to another function - RULE_GEN_FORM.reload_features_tree() - except Exception as e: - logger.warn("exception: %s", e) - - -class RuleGeneratorForm(idaapi.PluginForm): - def __init__(self): - super(RuleGeneratorForm, self).__init__() - self.title = "capa rule generator" - - self.parent = None - self.parent_items = {} - self.orig_colors = None - - self.hooks = Hooks() # dirty? - if self.hooks.hook(): - logger.info("UI notification hook installed successfully") - - def init_ui(self): - self.tree = QTreeWidget() - self.rule_text = QTextEdit() - self.rule_text.setMinimumWidth(350) - - self.reload_features_tree() - - button_reset = QtWidgets.QPushButton("&Reset") - button_reset.clicked.connect(self.reset) - - h_layout = QtWidgets.QHBoxLayout() - v_layout = QtWidgets.QVBoxLayout() - - h_layout.addWidget(self.tree) - h_layout.addWidget(self.rule_text) - - v_layout.addLayout(h_layout) - v_layout.addWidget(button_reset) - - self.parent.setLayout(v_layout) - - def reset(self): - plugin_helpers.reset_selection(self.tree) - plugin_helpers.reset_colors(self.orig_colors) - self.rule_text.setText("") - - def reload_features_tree(self): - self.reset() - self.tree.clear() - self.orig_colors = None - self.parent_items = {} - - features = self.get_features() - - if not features: - return - - feature_vas = set().union(*features.values()) - self.orig_colors = plugin_helpers.get_orig_color_feature_vas(feature_vas) - self.create_tree(features) - self.tree.update() - - def get_features(self): - # load like standalone tool - extractor = capa.features.extractors.ida.IdaFeatureExtractor() - f = idaapi.get_func(idaapi.get_screen_ea()) - if not f: - logger.info("function does not exist at 0x%x", idaapi.get_screen_ea()) - return - - return self.extract_function_features(f) - - def extract_function_features(self, f): - features = defaultdict(set) - for bb in idaapi.FlowChart(f, flags=idaapi.FC_PREDS): - for insn in capa.features.extractors.ida.helpers.get_instructions_in_range(bb.start_ea, bb.end_ea): - for feature, va in capa.features.extractors.ida.insn.extract_features(f, bb, insn): - features[feature].add(va) - for feature, va in capa.features.extractors.ida.basicblock.extract_features(f, bb): - features[feature].add(va) - return features - - def create_tree(self, features): - self.tree.setMinimumWidth(400) - # self.tree.setMinimumHeight(300) - self.tree.setHeaderLabels(["Feature", "Virtual Address", "Disassembly"]) - # auto resize columns - self.tree.header().setSectionResizeMode(QHeaderView.ResizeToContents) - self.tree.itemClicked.connect(self.on_item_clicked) - - # features sorted by location of first occurrence - # TODO fix characteristic features display and rule text - for feature, vas in sorted(features.items(), key=lambda k: sorted(k[1])): - # level 0 - if type(feature) not in self.parent_items: - self.parent_items[type(feature)] = plugin_helpers.add_child_item(self.tree, [feature.name.lower()]) - - # level 1 - if feature not in self.parent_items: - self.parent_items[feature] = plugin_helpers.add_child_item( - self.parent_items[type(feature)], [str(feature)] - ) - - # level n > 1 - if len(vas) > 1: - for va in sorted(vas): - plugin_helpers.add_child_item( - self.parent_items[feature], - [str(feature), "0x%X" % va, plugin_helpers.get_disasm_line(va)], - feature, - ) - else: - va = vas.pop() - self.parent_items[feature].setText(0, str(feature)) - self.parent_items[feature].setText(1, "0x%X" % va) - self.parent_items[feature].setText(2, plugin_helpers.get_disasm_line(va)) - self.parent_items[feature].setData(0, 0x100, feature) - - # @QtCore.pyqtSlot(QTreeWidgetItem, int) - def on_item_clicked(self, it, col): - # logger.debug('clicked %s, %s, %s', it, col, it.text(col)) - # jump to address - if col == 1 and it.text(col): - va = int(it.text(col), 0x10) - if va: - idc.jumpto(va) - - # highlight in disassembly - plugin_helpers.reset_colors(self.orig_colors) - selected = self.get_selected_items() - for va in selected.keys(): - idc.set_color(va, idc.CIC_ITEM, COLOR_HIGHLIGHT) - - self.update_rule_text() - - def update_rule_text(self): - features = self.get_selected_items().values() - rule = self.get_rule_from_features(features) - self.rule_text.setText(rule) - - def get_rule_from_features(self, features): - rule_parts = [] - counted = zip( - Counter(features).keys(), Counter(features).values() # equals to list(set(words)) - ) # counts the elements' frequency - - # single features - for k, v in filter(lambda t: t[1] == 1, counted): - # TODO args to hex if int - if k.name.lower() == "bytes": - # Convert raw bytes to uppercase hex representation (e.g., '12 34 56') - upper_hex_bytes = binascii.hexlify(args_to_str(k.args)).upper() - rule_value_str = "" - for i in range(0, len(upper_hex_bytes), 2): - rule_value_str += upper_hex_bytes[i : i + 2] + " " - r = " - %s: %s" % (k.name.lower(), rule_value_str) - else: - r = " - %s: %s" % (k.name.lower(), args_to_str(k.args)) - rule_parts.append(r) - - # counted features - for k, v in filter(lambda t: t[1] > 1, counted): - r = " - count(%s): %d" % (str(k), v) - rule_parts.append(r) - - rule_prefix = textwrap.dedent( - """ - rule: - meta: - name: - author: %s - scope: function - examples: - - %s:0x%X - features: - """ - % (AUTHOR_NAME, idc.retrieve_input_file_md5(), get_func_start(idc.here())) - ).strip() - return "%s\n%s" % (rule_prefix, "\n".join(sorted(rule_parts))) - - # TODO merge into capa_idautils, get feature data - def get_selected_items(self): - selected = {} - iterator = QtWidgets.QTreeWidgetItemIterator(self.tree, QtWidgets.QTreeWidgetItemIterator.Checked) - while iterator.value(): - item = iterator.value() - if item.text(1): - # logger.debug('selected %s, %s, %s', item.text(1), item.text(0), item.data(0, 0x100)) - selected[int(item.text(1), 0x10)] = item.data(0, 0x100) - iterator += 1 - return selected - - # ---------------------------------------------------------- - # IDA Plugin API - # ---------------------------------------------------------- - def OnCreate(self, form): - self.parent = self.FormToPyQtWidget(form) - self.init_ui() - - def Show(self): - return idaapi.PluginForm.Show( - self, self.title, options=(idaapi.PluginForm.WOPN_RESTORE | idaapi.PluginForm.WOPN_PERSIST) - ) - - def OnClose(self, form): - self.reset() - if self.hooks.unhook(): - logger.info("UI notification hook uninstalled successfully") - logger.info("RuleGeneratorForm closed") - - -def args_to_str(args): - a = [] - for arg in args: - if (isinstance(arg, int) or isinstance(arg, long)) and arg > 10: - a.append("0x%X" % arg) - else: - a.append(str(arg)) - return ",".join(a) - - -def main(): - logging.basicConfig(level=logging.INFO) - - global RULE_GEN_FORM - try: - # there is an instance, reload it - RULE_GEN_FORM - RULE_GEN_FORM.Close() - RULE_GEN_FORM = RuleGeneratorForm() - except Exception: - # there is no instance yet - RULE_GEN_FORM = RuleGeneratorForm() - - RULE_GEN_FORM.Show() - - -if __name__ == "__main__": - main()