adding support to run explorer as IDA plugin

README.md

@@ -146,7 +146,7 @@ rule:
 The [github.com/fireeye/capa-rules](https://github.com/fireeye/capa-rules) repository contains hundreds of standard library rules that are distributed with capa.
 Please learn to write rules and contribute new entries as you find interesting techniques in malware.
 
-If you use IDA Pro, then you use can use the [IDA Pro plugin for capa](./capa/ida/ida_capa_explorer.py).
+If you use IDA Pro, then you use can use the [IDA Pro plugin for capa](capa/ida/plugin/).
 This script adds new user interface elements to IDA, including an interactive tree view of rule matches and their locations within the current database.
 As you select the checkboxes, the plugin will highlight the addresses associated with the features.
 We use this plugin all the time to quickly jump to interesting parts of a program.

capa/ida/helpers/__init__.py

@@ -46,7 +46,7 @@ def is_supported_ida_version():
         logger.warning(
             "Your IDA Pro version is: %s. Supported versions are: %s." % (version, ", ".join(SUPPORTED_IDA_VERSIONS))
         )
-        capa.ida.helpers.inform_user_ida_ui(warning_msg)
+        # capa.ida.helpers.inform_user_ida_ui(warning_msg)
         return False
     return True
 
@@ -62,7 +62,7 @@ def is_supported_file_type():
         )
         logger.error(" If you don't know the input file type, you can try using the `file` utility to guess it.")
         logger.error("-" * 80)
-        inform_user_ida_ui("capa does not support the format of this file")
+        # inform_user_ida_ui("capa does not support the format of this file")
         return False
     return True
 

capa/ida/plugin/__init__.py

@@ -0,0 +1,66 @@
+# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at: [package root]/LICENSE.txt
+# 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.
+
+import logging
+
+import idaapi
+
+from capa.ida.helpers import is_supported_file_type, is_supported_ida_version
+from capa.ida.plugin.form import CapaExplorerForm
+
+logger = logging.getLogger("capa")
+
+
+class CapaExplorerPlugin(idaapi.plugin_t):
+
+    # Mandatory definitions
+    PLUGIN_NAME = "capa explorer"
+    PLUGIN_VERSION = "0.0.1"
+    PLUGIN_AUTHORS = ""
+
+    wanted_name = PLUGIN_NAME
+    comment = "IDA plugin for capa analysis framework"
+    version = ""
+    website = ""
+    help = ""
+    wanted_hotkey = ""
+    flags = 0
+
+    def __init__(self):
+        """ """
+        self.form = None
+
+    def init(self):
+        """
+        called when IDA is loading the plugin
+        """
+        logging.basicConfig(level=logging.INFO)
+
+        # check IDA version and database compat
+        if not is_supported_ida_version():
+            return idaapi.PLUGIN_SKIP
+        if not is_supported_file_type():
+            return idaapi.PLUGIN_SKIP
+
+        logger.info("plugin initialized.")
+
+        return idaapi.PLUGIN_KEEP
+
+    def term(self):
+        """
+        called when IDA is unloading the plugin
+        """
+        logger.info("plugin closed.")
+
+    def run(self, arg):
+        """
+        called when IDA is running the plugin as a script
+        """
+        self.form = CapaExplorerForm(self.PLUGIN_NAME, logger)
+        self.form.Show()
+        return True

capa/ida/plugin/form.py

@@ -8,84 +8,33 @@
 
 import os
 import json
-import logging
 import collections
 
-import idaapi
 from PyQt5 import QtGui, QtCore, QtWidgets
 
+import idaapi
+
 import capa.main
 import capa.rules
 import capa.ida.helpers
 import capa.render.utils as rutils
 import capa.features.extractors.ida
-from capa.ida.explorer.view import CapaExplorerQtreeView
-from capa.ida.explorer.model import CapaExplorerDataModel
-from capa.ida.explorer.proxy import CapaExplorerSortFilterProxyModel
 
-PLUGIN_NAME = "capa explorer"
-
-logger = logging.getLogger("capa")
-
-
-class CapaExplorerIdaHooks(idaapi.UI_Hooks):
-    def __init__(self, screen_ea_changed_hook, action_hooks):
-        """facilitate IDA UI hooks
-
-        @param screen_ea_changed_hook: function hook for IDA screen ea changed
-        @param action_hooks: dict of IDA action handles
-        """
-        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 = {}
-
-    def preprocess_action(self, name):
-        """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)
-
-        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:
-            return
-
-        self.process_action_handle(self.process_action_meta, post=True)
-        self.reset()
-
-    def screen_ea_changed(self, curr_ea, prev_ea):
-        """called after screen location is changed
-
-        @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()
+from capa.ida.plugin.view import CapaExplorerQtreeView
+from capa.ida.plugin.model import CapaExplorerDataModel
+from capa.ida.plugin.proxy import CapaExplorerSortFilterProxyModel
+from capa.ida.plugin.hooks import CapaExplorerIdaHooks
 
 
 class CapaExplorerForm(idaapi.PluginForm):
-    def __init__(self):
+    def __init__(self, name, logger):
         """ """
         super(CapaExplorerForm, self).__init__()
 
-        self.form_title = PLUGIN_NAME
-        self.file_loc = __file__
+        self.form_title = name
+        self.logger = logger
+
+        self.rule_path = ""
 
         self.parent = None
         self.ida_hooks = None
@@ -112,10 +61,11 @@ class CapaExplorerForm(idaapi.PluginForm):
 
         self.view_tree.reset()
 
-        logger.info("form created.")
+        self.logger.info("form created.")
 
     def Show(self):
         """ """
+        self.logger.info("form show.")
         return idaapi.PluginForm.Show(
             self, self.form_title, options=(idaapi.PluginForm.WOPN_TAB | idaapi.PluginForm.WCLS_CLOSE_LATER)
         )
@@ -124,8 +74,7 @@ class CapaExplorerForm(idaapi.PluginForm):
         """ form is closed """
         self.unload_ida_hooks()
         self.ida_reset()
-
-        logger.info("form closed.")
+        self.logger.info("form closed.")
 
     def load_interface(self):
         """ load user interface """
@@ -262,6 +211,7 @@ class CapaExplorerForm(idaapi.PluginForm):
         actions = (
             ("Reset view", "Reset plugin view", self.reset),
             ("Run analysis", "Run capa analysis on current database", self.reload),
+            ("Change rules directory...", "Select new rules directory", self.change_rules_dir),
             ("Export results...", "Export capa results as JSON file", self.export_json),
         )
 
@@ -276,9 +226,15 @@ class CapaExplorerForm(idaapi.PluginForm):
         if not self.doc:
             idaapi.info("No capa results to export.")
             return
+
         path = idaapi.ask_file(True, "*.json", "Choose file")
+
+        if not path:
+            return
+
         if os.path.exists(path) and 1 != idaapi.ask_yn(1, "File already exists. Overwrite?"):
             return
+
         with open(path, "wb") as export_file:
             export_file.write(
                 json.dumps(self.doc, sort_keys=True, cls=capa.render.CapaJsonObjectEncoder).encode("utf-8")
@@ -347,16 +303,29 @@ class CapaExplorerForm(idaapi.PluginForm):
 
     def load_capa_results(self):
         """ run capa analysis and render results in UI """
-        logger.info("-" * 80)
-        logger.info(" Using default embedded rules.")
-        logger.info(" ")
-        logger.info(" You can see the current default rule set here:")
-        logger.info("     https://github.com/fireeye/capa-rules")
-        logger.info("-" * 80)
+        if not self.rule_path:
+            rule_path = self.ask_user_directory()
+            if not rule_path:
+                capa.ida.helpers.inform_user_ida_ui("You must select a rules directory to use for analysis.")
+                self.logger.warning("no rules directory selected. nothing to do.")
+                return
+            self.rule_path = rule_path
 
-        rules_path = os.path.join(os.path.dirname(self.file_loc), "../..", "rules")
-        rules = capa.main.get_rules(rules_path)
-        rules = capa.rules.RuleSet(rules)
+        self.logger.info("-" * 80)
+        self.logger.info(" Using rules from %s." % self.rule_path)
+        self.logger.info(" ")
+        self.logger.info(" You can see the current default rule set here:")
+        self.logger.info("     https://github.com/fireeye/capa-rules")
+        self.logger.info("-" * 80)
+
+        try:
+            rules = capa.main.get_rules(self.rule_path)
+            rules = capa.rules.RuleSet(rules)
+        except (IOError, capa.rules.InvalidRule, capa.rules.InvalidRuleSet) as e:
+            capa.ida.helpers.inform_user_ida_ui("Failed to load rules from %s" % self.rule_path)
+            self.logger.error("failed to load rules from %s (%s)" % (self.rule_path, e))
+            self.rule_path = ""
+            return
 
         meta = capa.ida.helpers.collect_metadata()
 
@@ -369,24 +338,26 @@ class CapaExplorerForm(idaapi.PluginForm):
         # warn user binary file is loaded but still allow capa to process it
         # TODO: check specific architecture of binary files based on how user configured IDA processors
         if idaapi.get_file_type_name() == "Binary file":
-            logger.warning("-" * 80)
-            logger.warning(" Input file appears to be a binary file.")
-            logger.warning(" ")
-            logger.warning(
+            self.logger.warning("-" * 80)
+            self.logger.warning(" Input file appears to be a binary file.")
+            self.logger.warning(" ")
+            self.logger.warning(
                 " capa currently only supports analyzing binary files containing x86/AMD64 shellcode with IDA."
             )
-            logger.warning(
+            self.logger.warning(
                 " This means the results may be misleading or incomplete if the binary file loaded in IDA is not x86/AMD64."
             )
-            logger.warning(" If you don't know the input file type, you can try using the `file` utility to guess it.")
-            logger.warning("-" * 80)
+            self.logger.warning(
+                " If you don't know the input file type, you can try using the `file` utility to guess it."
+            )
+            self.logger.warning("-" * 80)
 
             capa.ida.helpers.inform_user_ida_ui("capa encountered warnings during analysis")
 
         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.logger.info("analysis completed.")
 
         self.doc = capa.render.convert_capabilities_to_result_document(meta, rules, capabilities)
 
@@ -396,7 +367,7 @@ class CapaExplorerForm(idaapi.PluginForm):
 
         self.set_view_tree_default_sort_order()
 
-        logger.info("render views completed.")
+        self.logger.info("render views completed.")
 
     def set_view_tree_default_sort_order(self):
         """ set capa tree view default sort order """
@@ -494,18 +465,18 @@ class CapaExplorerForm(idaapi.PluginForm):
         self.view_summary.setRowCount(0)
         self.load_capa_results()
 
-        logger.info("reload complete.")
-        idaapi.info("%s reload completed." % PLUGIN_NAME)
+        self.logger.info("reload complete.")
+        idaapi.info("%s reload completed." % self.form_title)
 
-    def reset(self):
+    def reset(self, checked):
         """reset UI elements
 
         e.g. checkboxes and IDA highlighting
         """
         self.ida_reset()
 
-        logger.info("reset completed.")
-        idaapi.info("%s reset completed." % PLUGIN_NAME)
+        self.logger.info("reset completed.")
+        idaapi.info("%s reset completed." % self.form_title)
 
     def slot_menu_bar_hovered(self, action):
         """display menu action tooltip
@@ -518,13 +489,13 @@ class CapaExplorerForm(idaapi.PluginForm):
             QtGui.QCursor.pos(), action.toolTip(), self.view_menu_bar, self.view_menu_bar.actionGeometry(action)
         )
 
-    def slot_checkbox_limit_by_changed(self):
+    def slot_checkbox_limit_by_changed(self, state):
         """slot activated if checkbox clicked
 
         if checked, configure function filter if screen location is located
         in function, otherwise clear filter
         """
-        if self.view_limit_results_by_function.isChecked():
+        if state == QtCore.Qt.Checked:
             self.limit_results_to_function(idaapi.get_func(idaapi.get_screen_ea()))
         else:
             self.model_proxy.reset_address_range_filter()
@@ -542,30 +513,16 @@ class CapaExplorerForm(idaapi.PluginForm):
             # if function not exists don't display any results (address should not be -1)
             self.model_proxy.add_address_range_filter(-1, -1)
 
+    def ask_user_directory(self):
+        """ create Qt dialog to ask user for a directory """
+        return str(QtWidgets.QFileDialog.getExistingDirectory(self.parent, "Select rules directory"))
 
-def main():
-    """ TODO: move to idaapi.plugin_t class """
-    logging.basicConfig(level=logging.INFO)
-
-    if not capa.ida.helpers.is_supported_ida_version():
-        return -1
-
-    if not capa.ida.helpers.is_supported_file_type():
-        return -1
-
-    global CAPA_EXPLORER_FORM
-
-    try:
-        # there is an instance, reload it
-        CAPA_EXPLORER_FORM
-        CAPA_EXPLORER_FORM.Close()
-        CAPA_EXPLORER_FORM = CapaExplorerForm()
-    except Exception:
-        # there is no instance yet
-        CAPA_EXPLORER_FORM = CapaExplorerForm()
-
-    CAPA_EXPLORER_FORM.Show()
-
-
-if __name__ == "__main__":
-    main()
+    def change_rules_dir(self):
+        """ allow user to change rules directory """
+        rule_path = self.ask_user_directory()
+        if not rule_path:
+            self.logger.warning("no rules directory selected. nothing to do.")
+            return
+        self.rule_path = rule_path
+        if 1 == idaapi.ask_yn(1, "Run analysis now?"):
+            self.reload()

capa/ida/plugin/hooks.py

@@ -0,0 +1,60 @@
+# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at: [package root]/LICENSE.txt
+# 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.
+
+import idaapi
+
+
+class CapaExplorerIdaHooks(idaapi.UI_Hooks):
+    def __init__(self, screen_ea_changed_hook, action_hooks):
+        """facilitate IDA UI hooks
+
+        @param screen_ea_changed_hook: function hook for IDA screen ea changed
+        @param action_hooks: dict of IDA action handles
+        """
+        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 = {}
+
+    def preprocess_action(self, name):
+        """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)
+
+        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:
+            return
+
+        self.process_action_handle(self.process_action_meta, post=True)
+        self.reset()
+
+    def screen_ea_changed(self, curr_ea, prev_ea):
+        """called after screen location is changed
+
+        @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()

capa/ida/plugin/model.py

@@ -9,13 +9,12 @@
 from collections import deque
 
 import idc
-import six
 import idaapi
-from PyQt5 import Qt, QtGui, QtCore
+from PyQt5 import QtGui, QtCore
 
 import capa.ida.helpers
 import capa.render.utils as rutils
-from capa.ida.explorer.item import (
+from capa.ida.plugin.item import (
     CapaExplorerDataItem,
     CapaExplorerRuleItem,
     CapaExplorerBlockItem,

capa/ida/plugin/proxy.py

@@ -8,7 +8,7 @@
 
 from PyQt5 import QtCore
 
-from capa.ida.explorer.model import CapaExplorerDataModel
+from capa.ida.plugin.model import CapaExplorerDataModel
 
 
 class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel):

capa/ida/plugin/view.py

@@ -7,11 +7,10 @@
 # See the License for the specific language governing permissions and limitations under the License.
 
 import idc
-import idaapi
-from PyQt5 import QtGui, QtCore, QtWidgets
+from PyQt5 import QtCore, QtWidgets
 
-from capa.ida.explorer.item import CapaExplorerRuleItem, CapaExplorerFunctionItem
-from capa.ida.explorer.model import CapaExplorerDataModel
+from capa.ida.plugin.item import CapaExplorerFunctionItem
+from capa.ida.plugin.model import CapaExplorerDataModel
 
 
 class CapaExplorerQtreeView(QtWidgets.QTreeView):

capa/ida/plugin_helpers.py

@@ -1,99 +0,0 @@
-# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at: [package root]/LICENSE.txt
-# 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.
-
-import os
-import logging
-
-import idc
-import idaapi
-from PyQt5.QtCore import Qt
-from PyQt5.QtWidgets import QTreeWidgetItem, QTreeWidgetItemIterator
-
-CAPA_EXTENSION = ".capas"
-
-
-logger = logging.getLogger("capa_ida")
-
-
-def get_input_file(freeze=True):
-    """
-    get input file path
-
-        freeze (bool): if True, get freeze file if it exists
-    """
-    # try original file in same directory as idb/i64 without idb/i64 file extension
-    input_file = idc.get_idb_path()[:-4]
-
-    if freeze:
-        # use frozen file if it exists
-        freeze_file_cand = "%s%s" % (input_file, CAPA_EXTENSION)
-        if os.path.isfile(freeze_file_cand):
-            return freeze_file_cand
-
-    if not os.path.isfile(input_file):
-        # TM naming
-        input_file = "%s.mal_" % idc.get_idb_path()[:-4]
-        if not os.path.isfile(input_file):
-            input_file = idaapi.ask_file(0, "*.*", "Please specify input file.")
-    if not input_file:
-        raise ValueError("could not find input file")
-    return input_file
-
-
-def get_orig_color_feature_vas(vas):
-    orig_colors = {}
-    for va in vas:
-        orig_colors[va] = idc.get_color(va, idc.CIC_ITEM)
-    return orig_colors
-
-
-def reset_colors(orig_colors):
-    if orig_colors:
-        for va, color in orig_colors.iteritems():
-            idc.set_color(va, idc.CIC_ITEM, orig_colors[va])
-
-
-def reset_selection(tree):
-    iterator = QTreeWidgetItemIterator(tree, QTreeWidgetItemIterator.Checked)
-    while iterator.value():
-        item = iterator.value()
-        item.setCheckState(0, Qt.Unchecked)  # column, state
-        iterator += 1
-
-
-def get_disasm_line(va):
-    return idc.generate_disasm_line(va, idc.GENDSM_FORCE_CODE)
-
-
-def get_selected_items(tree, skip_level_1=False):
-    selected = []
-    iterator = QTreeWidgetItemIterator(tree, QTreeWidgetItemIterator.Checked)
-    while iterator.value():
-        item = iterator.value()
-        if skip_level_1:
-            # hacky way to check if item is at level 1, if so, skip
-            # alternative, check if text in disasm column
-            if item.parent() and item.parent().parent() is None:
-                iterator += 1
-                continue
-        if item.text(1):
-            # logger.debug('selected %s, %s', item.text(0), item.text(1))
-            selected.append(int(item.text(1), 0x10))
-        iterator += 1
-    return selected
-
-
-def add_child_item(parent, values, feature=None):
-    child = QTreeWidgetItem(parent)
-    child.setFlags(child.flags() | Qt.ItemIsTristate | Qt.ItemIsUserCheckable)
-    for i, v in enumerate(values):
-        child.setText(i, v)
-        if feature:
-            child.setData(0, 0x100, feature)
-        child.setCheckState(0, Qt.Unchecked)
-    return child

capa_plugin_ida.py

@@ -0,0 +1,14 @@
+# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at: [package root]/LICENSE.txt
+# 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 capa.ida.plugin import CapaExplorerPlugin
+
+
+def PLUGIN_ENTRY():
+    """ Mandatory entry point for IDAPython plugins """
+    return CapaExplorerPlugin()

doc/usage.md

@@ -22,7 +22,7 @@ IDA's analysis is generally a bit faster and more thorough than vivisect's, so y
 When run under IDA, capa supports both Python 2 and Python 3 interpreters.
 If you encounter issues with your specific setup, please open a new [Issue](https://github.com/fireeye/capa/issues).
 
-Additionally, capa comes with an IDA Pro plugin located in the `capa/ida` directory: the explorer.
+Additionally, capa comes with an IDA Pro plugin located in the `capa/ida/plugin` directory: the explorer.
 
 #### capa explorer
 The capa explorer allows you to interactively display and browse capabilities capa identified in a binary.
@@ -31,10 +31,8 @@ We like to use capa to help find the most interesting parts of a program, such a
 
 ![capa explorer](img/capa_explorer.png)
 
-To install the plugin, you'll need to be running IDA Pro 7.4 or 7.5 with either Python 2 or Python 3.
-Next make sure pip commands are run using the Python install that is configured for your IDA install:
-
-  1. Only if running Python 2.7, run command `$ pip install https://github.com/williballenthin/vivisect/zipball/master`
-  2. Run `$ pip install .` from capa root directory
-  3. Open IDA and navigate to `File > Script file…` or `Alt+F7`
-  4. Navigate to `<capa_install_dir>\capa\ida\` and choose `ida_capa_explorer.py`
+The plugin currently supports IDA Pro 7.1 through 7.5 with either Python 2 or Python 3. To use the plugin, install capa
+by following method 2 or 3 from the [installation guide](doc/installation.md) and copy [capa_plugin_ida.py](capa_plugin_ida.py)
+to the plugins directory of your IDA Pro installation. Following these steps you can run capa explorer in IDA Pro by navigating 
+to `Edit > Plugins > capa explorer`. The plugin will prompt you to select a rules directory to use for analysis. You can 
+use the [default rule set](https://github.com/fireeye/capa-rules/) or point the plugin to your own directory of rules.