From f9451feb1804b92716f99b022857ef8962fae0bb Mon Sep 17 00:00:00 2001 From: Michael Hunhoff Date: Tue, 8 Sep 2020 12:26:20 -0600 Subject: [PATCH] changes to plugin function-level documentation --- capa/ida/plugin/__init__.py | 17 ++- capa/ida/plugin/capa_plugin_ida.py | 5 +- capa/ida/plugin/form.py | 126 ++++++++++++++-------- capa/ida/plugin/hooks.py | 4 +- capa/ida/plugin/item.py | 166 +++++++++++++++++++++-------- capa/ida/plugin/model.py | 134 +++++++++-------------- capa/ida/plugin/proxy.py | 58 ++++++---- capa/ida/plugin/view.py | 75 ++++++------- 8 files changed, 338 insertions(+), 247 deletions(-) diff --git a/capa/ida/plugin/__init__.py b/capa/ida/plugin/__init__.py index 8b651509..130d4a68 100644 --- a/capa/ida/plugin/__init__.py +++ b/capa/ida/plugin/__init__.py @@ -34,13 +34,11 @@ class CapaExplorerPlugin(idaapi.plugin_t): flags = 0 def __init__(self): - """ """ + """initialize plugin""" self.form = None def init(self): - """ - called when IDA is loading the plugin - """ + """called when IDA is loading the plugin""" logging.basicConfig(level=logging.INFO) # check IDA version and database compatibility @@ -51,18 +49,15 @@ class CapaExplorerPlugin(idaapi.plugin_t): logger.debug("plugin initialized") - return idaapi.PLUGIN_KEEP + # plugin is good, but don't keep us in memory + return idaapi.PLUGIN_OK def term(self): - """ - called when IDA is unloading the plugin - """ + """called when IDA is unloading the plugin""" logger.debug("plugin terminated") def run(self, arg): - """ - called when IDA is running the plugin as a script - """ + """called when IDA is running the plugin as a script""" self.form = CapaExplorerForm(self.PLUGIN_NAME) self.form.Show() return True diff --git a/capa/ida/plugin/capa_plugin_ida.py b/capa/ida/plugin/capa_plugin_ida.py index c0854c0f..f8b4ae72 100644 --- a/capa/ida/plugin/capa_plugin_ida.py +++ b/capa/ida/plugin/capa_plugin_ida.py @@ -10,5 +10,8 @@ from capa.ida.plugin import CapaExplorerPlugin def PLUGIN_ENTRY(): - """ Mandatory entry point for IDAPython plugins """ + """mandatory entry point for IDAPython plugins + + copy this script to your IDA plugins directory and start the plugin by navigating to Edit > Plugins in IDA Pro + """ return CapaExplorerPlugin() diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index 201c1c16..1d325281 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -31,8 +31,10 @@ settings = ida_settings.IDASettings("capa") class CapaExplorerForm(idaapi.PluginForm): + """form element for plugin interface""" + def __init__(self, name): - """ """ + """initialize form elements""" super(CapaExplorerForm, self).__init__() self.form_title = name @@ -56,9 +58,11 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_menu_bar = None def OnCreate(self, form): - """ """ + """called when plugin form is created""" self.parent = self.FormToPyQtWidget(form) self.parent.setWindowIcon(QICON) + + # load interface elements self.load_interface() self.load_capa_results() self.load_ida_hooks() @@ -68,20 +72,23 @@ class CapaExplorerForm(idaapi.PluginForm): logger.debug("form created") def Show(self): - """ """ + """creates form if not already create, else brings plugin to front""" logger.debug("form show") 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 """ + """called when form is closed + + ensure any plugin modifications (e.g. hooks and UI changes) are reset before the plugin is closed + """ self.unload_ida_hooks() self.ida_reset() logger.debug("form closed") def load_interface(self): - """ load user interface """ + """load user interface""" # load models self.model_data = CapaExplorerDataModel() @@ -113,17 +120,17 @@ class CapaExplorerForm(idaapi.PluginForm): self.load_view_parent() def load_view_tabs(self): - """ load tabs """ + """load tabs""" tabs = QtWidgets.QTabWidget() self.view_tabs = tabs def load_view_menu_bar(self): - """ load menu bar """ + """load menu bar""" bar = QtWidgets.QMenuBar() self.view_menu_bar = bar def load_view_attack(self): - """ load MITRE ATT&CK table """ + """load MITRE ATT&CK table""" table_headers = [ "ATT&CK Tactic", "ATT&CK Technique ", @@ -145,7 +152,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_attack = table def load_view_checkbox_limit_by(self): - """ load limit results by function checkbox """ + """load limit results by function checkbox""" check = QtWidgets.QCheckBox("Limit results to current function") check.setChecked(False) check.stateChanged.connect(self.slot_checkbox_limit_by_changed) @@ -153,7 +160,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_limit_results_by_function = check def load_view_search_bar(self): - """ load the search bar control """ + """load the search bar control""" line = QtWidgets.QLineEdit() line.setPlaceholderText("search...") line.textChanged.connect(self.search_model_proxy.set_query) @@ -161,7 +168,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_search_bar = line def load_view_parent(self): - """ load view parent """ + """load view parent""" layout = QtWidgets.QVBoxLayout() layout.addWidget(self.view_tabs) @@ -170,7 +177,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.parent.setLayout(layout) def load_view_tree_tab(self): - """ load capa tree tab view """ + """load tree view tab""" layout = QtWidgets.QVBoxLayout() layout.addWidget(self.view_limit_results_by_function) layout.addWidget(self.view_search_bar) @@ -182,7 +189,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_tabs.addTab(tab, "Tree View") def load_view_attack_tab(self): - """ load MITRE ATT&CK tab view """ + """load MITRE ATT&CK view tab""" layout = QtWidgets.QVBoxLayout() layout.addWidget(self.view_attack) @@ -192,6 +199,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_tabs.addTab(tab, "MITRE") def load_file_menu(self): + """load file menu controls""" actions = ( ("Rerun analysis", "Rerun capa analysis on current database", self.reload), ("Export results...", "Export capa results as JSON file", self.export_json), @@ -199,15 +207,21 @@ class CapaExplorerForm(idaapi.PluginForm): self.load_menu("File", actions) def load_rules_menu(self): + """load rules menu controls""" actions = (("Change rules directory...", "Select new rules directory", self.change_rules_dir),) self.load_menu("Rules", actions) def load_view_menu(self): + """load view menu controls""" actions = (("Reset view", "Reset plugin view", self.reset),) self.load_menu("View", actions) def load_menu(self, title, actions): - """ load menu actions """ + """load menu actions + + @param title: menu name displayed in UI + @param actions: tuple of tuples containing action name, tooltip, and slot function + """ menu = self.view_menu_bar.addMenu(title) for (name, _, handle) in actions: action = QtWidgets.QAction(name, self.parent) @@ -215,16 +229,18 @@ class CapaExplorerForm(idaapi.PluginForm): menu.addAction(action) def export_json(self): - """ export capa results as JSON file """ + """export capa results as JSON file""" if not self.doc: idaapi.info("No capa results to export.") return path = idaapi.ask_file(True, "*.json", "Choose file") + # user cancelled, entered blank input, etc. if not path: return + # check file exists, ask to override if os.path.exists(path) and 1 != idaapi.ask_yn(1, "File already exists. Overwrite?"): return @@ -234,7 +250,9 @@ class CapaExplorerForm(idaapi.PluginForm): ) def load_ida_hooks(self): - """ load IDA Pro UI hooks """ + """load IDA UI hooks""" + + # map named action (defined in idagui.cfg) to Python function action_hooks = { "MakeName": self.ida_hook_rename, "EditFunction": self.ida_hook_rename, @@ -245,18 +263,20 @@ class CapaExplorerForm(idaapi.PluginForm): self.ida_hooks.hook() def unload_ida_hooks(self): - """ unload IDA Pro UI hooks """ + """unload IDA Pro UI hooks + + must be called before plugin is completely destroyed + """ if self.ida_hooks: self.ida_hooks.unhook() def ida_hook_rename(self, meta, post=False): - """hook for IDA rename action + """function hook for IDA "MakeName" and "EditFunction" actions - called twice, once before action and once after - action completes + called twice, once before action and once after action completes - @param meta: metadata cache - @param post: indicates pre or post action + @param meta: dict of key/value pairs set when action first called (may be empty) + @param post: False if action first call, True if action second call """ location = idaapi.get_screen_ea() if not location or not capa.ida.helpers.is_func_start(location): @@ -272,9 +292,10 @@ class CapaExplorerForm(idaapi.PluginForm): meta["prev_name"] = curr_name def ida_hook_screen_ea_changed(self, widget, new_ea, old_ea): - """hook for IDA screen ea changed + """function hook for IDA "screen ea changed" action - this hook is currently only relevant for limiting results displayed in the UI + called twice, once before action and once after action completes. this hook is currently only relevant + for limiting results displayed in the UI @param widget: IDA widget type @param new_ea: destination ea @@ -296,20 +317,21 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_tree.resize_columns_to_content() def ida_hook_rebase(self, meta, post=False): - """hook for IDA rebase action + """function hook for IDA "RebaseProgram" action - called twice, once before action and once after - action completes + called twice, once before action and once after action completes - @param meta: metadata cache - @param post: indicates pre or post action + @param meta: dict of key/value pairs set when action first called (may be empty) + @param post: False if action first call, True if action second call """ if post: capa.ida.helpers.inform_user_ida_ui("Running capa analysis again after rebase") self.reload() def load_capa_results(self): - """ run capa analysis and render results in UI """ + """run capa analysis and render results in UI""" + + # resolve rules directory - check self and settings first, then ask user if not self.rule_path: if "rule_path" in settings: self.rule_path = settings["rule_path"] @@ -370,6 +392,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.doc = capa.render.convert_capabilities_to_result_document(meta, rules, capabilities) + # render views self.model_data.render_capa_doc(self.doc) self.render_capa_doc_mitre_summary() @@ -378,11 +401,11 @@ class CapaExplorerForm(idaapi.PluginForm): logger.debug("render views completed.") def set_view_tree_default_sort_order(self): - """ set capa tree view default sort order """ + """set tree view default sort order""" self.view_tree.sortByColumn(CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION, QtCore.Qt.AscendingOrder) def render_capa_doc_mitre_summary(self): - """ render capa MITRE ATT&CK results """ + """render MITRE ATT&CK results""" tactics = collections.defaultdict(set) for rule in rutils.capability_rules(self.doc): @@ -419,29 +442,32 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_attack.setRowCount(max(len(column_one), len(column_two))) - for row, value in enumerate(column_one): + 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): + 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): - """ create new table header item with default style """ + """create new table header item with our style + + @param text: header text to display + """ 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 UI """ + """reset plugin UI + + called when user selects plugin reset from menu + """ self.model_data.reset() self.view_tree.reset() self.view_limit_results_by_function.setChecked(False) @@ -449,7 +475,10 @@ class CapaExplorerForm(idaapi.PluginForm): self.set_view_tree_default_sort_order() def reload(self): - """ reload views and re-run capa analysis """ + """re-run capa analysis and reload UI controls + + called when user selects plugin reload from menu + """ self.ida_reset() self.range_model_proxy.invalidate() self.search_model_proxy.invalidate() @@ -483,8 +512,9 @@ class CapaExplorerForm(idaapi.PluginForm): 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 checked, configure function filter if screen location is located in function, otherwise clear filter + + @param state: checked state """ if state == QtCore.Qt.Checked: self.limit_results_to_function(idaapi.get_func(idaapi.get_screen_ea())) @@ -496,20 +526,26 @@ class CapaExplorerForm(idaapi.PluginForm): def limit_results_to_function(self, f): """add filter to limit results to current function + adds new address range filter to include function bounds, allowing basic blocks matched within a function + to be included in the results + @param f: (IDA func_t) """ if f: self.range_model_proxy.add_address_range_filter(f.start_ea, f.end_ea) else: - # if function not exists don't display any results (address should not be -1) + # if function not exists don't display any results (assume address never -1) self.range_model_proxy.add_address_range_filter(-1, -1) def ask_user_directory(self): - """ create Qt dialog to ask user for a directory """ + """create Qt dialog to ask user for a directory""" return str(QtWidgets.QFileDialog.getExistingDirectory(self.parent, "Select rules directory", self.rule_path)) def change_rules_dir(self): - """ allow user to change rules directory """ + """allow user to change rules directory + + user selection stored in settings for future runs + """ rule_path = self.ask_user_directory() if not rule_path: logger.warning("no rules directory selected. nothing to do.") diff --git a/capa/ida/plugin/hooks.py b/capa/ida/plugin/hooks.py index 4ad688cc..04e540c0 100644 --- a/capa/ida/plugin/hooks.py +++ b/capa/ida/plugin/hooks.py @@ -39,7 +39,7 @@ class CapaExplorerIdaHooks(idaapi.UI_Hooks): return 0 def postprocess_action(self): - """ called after action completed """ + """called after action completed""" if not self.process_action_handle: return @@ -55,6 +55,6 @@ class CapaExplorerIdaHooks(idaapi.UI_Hooks): self.screen_ea_changed_hook(idaapi.get_current_widget(), curr_ea, prev_ea) def reset(self): - """ reset internal state """ + """reset internal state""" self.process_action_handle = None self.process_action_meta.clear() diff --git a/capa/ida/plugin/item.py b/capa/ida/plugin/item.py index 46111633..b63886fc 100644 --- a/capa/ida/plugin/item.py +++ b/capa/ida/plugin/item.py @@ -28,20 +28,21 @@ def info_to_name(display): def location_to_hex(location): - """ convert location to hex for display """ + """convert location to hex for display""" return "%08X" % location class CapaExplorerDataItem(object): - """ store data for CapaExplorerDataModel """ + """store data for CapaExplorerDataModel""" def __init__(self, parent, data): - """ """ + """initialize item""" self.pred = parent self._data = data self.children = [] self._checked = False + # default state for item self.flags = ( QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable @@ -53,117 +54,146 @@ class CapaExplorerDataItem(object): self.pred.appendChild(self) def setIsEditable(self, isEditable=False): - """ modify item flags to be editable or not """ + """modify item editable flags + + @param isEditable: True, can edit, False cannot edit + """ 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 + + @param checked: True, item checked, False item not 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 a new child to specified item - @param item: CapaExplorerDataItem* + @param item: CapaExplorerDataItem """ self.children.append(item) def child(self, row): """get child row - @param row: TODO + @param row: row number """ return self.children[row] def childCount(self): - """ get child count """ + """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 + + @param: column number + """ try: return self._data[column] except IndexError: return None def parent(self): - """ get parent """ + """get parent""" return self.pred def row(self): - """ get row location """ + """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 + + @param column: column number + @value: value to set (assume str) + """ self._data[column] = value def children(self): - """ yield children """ + """yield children""" for child in self.children: yield child def removeChildren(self): - """ remove children from node """ + """remove children""" del self.children[:] def __str__(self): - """ get string representation of columns """ + """get string representation of columns + + used for copy-n-paste operations + """ return " ".join([data for data in self._data if data]) @property def info(self): - """ return data stored in information column """ + """return data stored in information column""" return self._data[0] @property def location(self): - """ return data stored in location column """ + """return data stored in location column""" try: + # address stored as str, convert to int before return return int(self._data[1], 16) except ValueError: return None @property def details(self): - """ return data stored in details column """ + """return data stored in details column""" return self._data[2] class CapaExplorerRuleItem(CapaExplorerDataItem): - """ store data relevant to capa function result """ + """store data for rule result""" fmt = "%s (%d matches)" def __init__(self, parent, name, namespace, count, source): - """ """ + """initialize item + + @param parent: parent node + @param name: rule name + @param namespace: rule namespace + @param count: number of match for this rule + @param source: rule source (tooltip) + """ display = self.fmt % (name, count) if count > 1 else name super(CapaExplorerRuleItem, self).__init__(parent, [display, "", namespace]) self._source = source @property def source(self): - """ return rule contents for display """ + """return rule source to display (tooltip)""" return self._source class CapaExplorerRuleMatchItem(CapaExplorerDataItem): - """ store data relevant to capa function match result """ + """store data for rule match""" def __init__(self, parent, display, source=""): - """ """ + """initialize item + + @param parent: parent node + @param display: text to display in UI + @param source: rule match source to display (tooltip) + """ super(CapaExplorerRuleMatchItem, self).__init__(parent, [display, "", ""]) self._source = source @@ -174,82 +204,125 @@ class CapaExplorerRuleMatchItem(CapaExplorerDataItem): class CapaExplorerFunctionItem(CapaExplorerDataItem): - """ store data relevant to capa function result """ + """store data for function match""" fmt = "function(%s)" def __init__(self, parent, location): - """ """ + """initialize item + + @param parent: parent node + @param location: virtual address of function as seen by IDA + """ super(CapaExplorerFunctionItem, self).__init__( parent, [self.fmt % idaapi.get_name(location), location_to_hex(location), ""] ) @property def info(self): - """ """ + """return function name""" info = super(CapaExplorerFunctionItem, self).info display = info_to_name(info) return display if display else info @info.setter def info(self, display): - """ """ + """set function name + + called when user changes function name in plugin UI + + @param display: new function name to display + """ self._data[0] = self.fmt % display class CapaExplorerSubscopeItem(CapaExplorerDataItem): - """ store data relevant to subscope """ + """store data for subscope match""" fmt = "subscope(%s)" def __init__(self, parent, scope): - """ """ + """initialize item + + @param parent: parent node + @param scope: subscope name + """ super(CapaExplorerSubscopeItem, self).__init__(parent, [self.fmt % scope, "", ""]) class CapaExplorerBlockItem(CapaExplorerDataItem): - """ store data relevant to capa basic block result """ + """store data for basic block match""" fmt = "basic block(loc_%08X)" def __init__(self, parent, location): - """ """ + """initialize item + + @param parent: parent node + @param location: virtual address of basic block as seen by IDA + """ super(CapaExplorerBlockItem, self).__init__(parent, [self.fmt % location, location_to_hex(location), ""]) class CapaExplorerDefaultItem(CapaExplorerDataItem): - """ store data relevant to capa default result """ + """store data for default match e.g. statement (and, or)""" def __init__(self, parent, display, details="", location=None): - """ """ + """initialize item + + @param parent: parent node + @param display: text to display in UI + @param details: text to display in details section of UI + @param location: virtual address as seen by IDA + """ 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 for feature match""" def __init__(self, parent, display, location="", details=""): - """ """ + """initialize item + + @param parent: parent node + @param display: text to display in UI + @param details: text to display in details section of UI + @param location: virtual address as seen by IDA + """ location = location_to_hex(location) if location else "" super(CapaExplorerFeatureItem, self).__init__(parent, [display, location, details]) class CapaExplorerInstructionViewItem(CapaExplorerFeatureItem): - """ store data relevant to an instruction preview """ + """store data for instruction match""" def __init__(self, parent, display, location): - """ """ + """initialize item + + details section shows disassembly view for match + + @param parent: parent node + @param display: text to display in UI + @param location: virtual address as seen by IDA + """ 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): - """ store data relevant to byte preview """ + """store data for byte match""" def __init__(self, parent, display, location): - """ """ + """initialize item + + details section shows byte preview for match + + @param parent: parent node + @param display: text to display in UI + @param location: virtual address as seen by IDA + """ byte_snap = idaapi.get_bytes(location, 32) if byte_snap: @@ -266,9 +339,14 @@ class CapaExplorerByteViewItem(CapaExplorerFeatureItem): class CapaExplorerStringViewItem(CapaExplorerFeatureItem): - """ store data relevant to string preview """ + """store data for string match""" def __init__(self, parent, display, location): - """ """ + """initialize item + + @param parent: parent node + @param display: text to display in UI + @param location: virtual address as seen by IDA + """ super(CapaExplorerStringViewItem, self).__init__(parent, display, location=location) self.ida_highlight = idc.get_color(location, idc.CIC_ITEM) diff --git a/capa/ida/plugin/model.py b/capa/ida/plugin/model.py index 0dadde19..93ecd3b7 100644 --- a/capa/ida/plugin/model.py +++ b/capa/ida/plugin/model.py @@ -33,7 +33,7 @@ DEFAULT_HIGHLIGHT = 0xD096FF class CapaExplorerDataModel(QtCore.QAbstractItemModel): - """ """ + """model for displaying hierarchical results return by capa""" COLUMN_INDEX_RULE_INFORMATION = 0 COLUMN_INDEX_VIRTUAL_ADDRESS = 1 @@ -42,14 +42,16 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): COLUMN_COUNT = 3 def __init__(self, parent=None): - """ """ + """initialize model""" super(CapaExplorerDataModel, self).__init__(parent) + # root node does not have parent, contains header columns self.root_node = CapaExplorerDataItem(None, ["Rule Information", "Address", "Details"]) def reset(self): - """ """ - # reset checkboxes and color highlights - # TODO: make less hacky + """reset UI elements (e.g. checkboxes, IDA color highlights) + + called when view wants to reset UI display + """ 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): @@ -58,15 +60,18 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): self.dataChanged.emit(model_index, model_index) def clear(self): - """ """ + """clear model data + + called when view wants to clear UI display + """ 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 + """return number of columns for the children of the given parent - @param model_index: QModelIndex* + @param model_index: QModelIndex @retval column count """ @@ -76,9 +81,11 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): 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 + """return data stored at given index by display role - @param model_index: QModelIndex* + this function is used to control UI elements (e.g. text font, color, etc.) based on column, item type, etc. + + @param model_index: QModelIndex @param role: QtCore.Qt.* @retval data to be displayed @@ -150,9 +157,9 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): return None def flags(self, model_index): - """get item flags for given index + """return item flags for given index - @param model_index: QModelIndex* + @param model_index: QModelIndex @retval QtCore.Qt.ItemFlags """ @@ -162,13 +169,13 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): 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 + """return 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() + @retval header data """ if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole: return self.root_node.data(section) @@ -176,13 +183,13 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): 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 + """return index of the item by row, column, and parent index - @param row: int - @param column: int - @param parent: QModelIndex* + @param row: item row + @param column: item column + @param parent: QModelIndex of parent - @retval QModelIndex* + @retval QModelIndex of item """ if not self.hasIndex(row, column, parent): return QtCore.QModelIndex() @@ -200,13 +207,13 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): return QtCore.QModelIndex() def parent(self, model_index): - """get parent of the model item with the given index + """return parent index by child index - if the item has no parent, an invalid QModelIndex* is returned + if the item has no parent, an invalid QModelIndex is returned - @param model_index: QModelIndex* + @param model_index: QModelIndex of child - @retval QModelIndex* + @retval QModelIndex of parent """ if not model_index.isValid(): return QtCore.QModelIndex() @@ -222,10 +229,10 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): 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 + @param model_index: QModelIndex of starting item + @param ignore_root: True, do not yield root index, False yield root index - @retval yield QModelIndex* + @retval yield QModelIndex """ visited = set() stack = deque((model_index,)) @@ -247,10 +254,10 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): stack.append(child_index.child(idx, 0)) def reset_ida_highlighting(self, item, checked): - """reset IDA highlight for an item + """reset IDA highlight for item - @param item: capa explorer item - @param checked: indicates item is or not checked + @param item: CapaExplorerDataItem + @param checked: True, item checked, False item not checked """ if not isinstance( item, (CapaExplorerStringViewItem, CapaExplorerInstructionViewItem, CapaExplorerByteViewItem) @@ -274,13 +281,11 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): 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 + """set data at index by role - @param model_index: QModelIndex* - @param value: QVariant* + @param model_index: QModelIndex of item + @param value: value to set @param role: QtCore.Qt.EditRole - - @retval True/False """ if not model_index.isValid(): return False @@ -315,12 +320,11 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): return False def rowCount(self, model_index): - """get the number of rows under the given parent + """return number of rows under item by index - when the parent is valid it means that is returning the number of - children of parent + when the parent is valid it means that is returning the number of children of parent - @param model_index: QModelIndex* + @param model_index: QModelIndex @retval row count """ @@ -340,11 +344,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): @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" - }, + @param doc: result doc """ if statement["type"] in ("and", "or", "optional"): display = statement["type"] @@ -398,24 +398,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): @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 - } - }, + @param doc: result doc """ if not match["success"]: # TODO: display failed branches at some point? Help with debugging rules? @@ -475,15 +458,6 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): """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", ""): @@ -500,13 +474,6 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): @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) @@ -535,14 +502,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): @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" - } + @param display: text to display in plugin UI """ # special handling for characteristic pending type if feature["type"] == "characteristic": @@ -598,7 +558,9 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): 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 + called when user updates function name using plugin UI + + @param old_name: old function name @param new_name: new function name """ # create empty root index for search diff --git a/capa/ida/plugin/proxy.py b/capa/ida/plugin/proxy.py index af9df7c1..b33a6f1b 100644 --- a/capa/ida/plugin/proxy.py +++ b/capa/ida/plugin/proxy.py @@ -13,20 +13,25 @@ from capa.ida.plugin.model import CapaExplorerDataModel class CapaExplorerRangeProxyModel(QtCore.QSortFilterProxyModel): - def __init__(self, parent=None): - """ """ - super(CapaExplorerRangeProxyModel, self).__init__(parent) + """filter results based on virtual address range as seen by IDA + implements filtering for "limit results by current function" checkbox in plugin UI + + minimum and maximum virtual addresses are used to filter results to a specific address range. this allows + basic blocks to be included when limiting results to a specific function + """ + + def __init__(self, parent=None): + """initialize proxy filter""" + super(CapaExplorerRangeProxyModel, self).__init__(parent) self.min_ea = None self.max_ea = None def lessThan(self, left, right): - """true if the value of the left item is less than value of right item + """return True if left item is less than right item, else False - @param left: QModelIndex* - @param right: QModelIndex* - - @retval True/False + @param left: QModelIndex of left + @param right: QModelIndex of right """ ldata = left.internalPointer().data(left.column()) rdata = right.internalPointer().data(right.column()) @@ -44,13 +49,11 @@ class CapaExplorerRangeProxyModel(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 - should be included in the model; otherwise returns false + """return true if the item in the row indicated by the given row and parent should be included in the model; + otherwise return false - @param row: int - @param parent: QModelIndex* - - @retval True/False + @param row: row number + @param parent: QModelIndex of parent """ if self.filter_accepts_row_self(row, parent): return True @@ -67,7 +70,11 @@ class CapaExplorerRangeProxyModel(QtCore.QSortFilterProxyModel): return False def index_has_accepted_children(self, row, parent): - """ """ + """return True if parent has one or more children that match filter, else False + + @param row: row number + @param parent: QModelIndex of parent + """ model_index = self.sourceModel().index(row, 0, parent) if model_index.isValid(): @@ -80,7 +87,11 @@ class CapaExplorerRangeProxyModel(QtCore.QSortFilterProxyModel): return False def filter_accepts_row_self(self, row, parent): - """ """ + """return True if filter accepts row, else False + + @param row: row number + @param parent: QModelIndex of parent + """ # filter not set if self.min_ea is None and self.max_ea is None: return True @@ -88,9 +99,11 @@ class CapaExplorerRangeProxyModel(QtCore.QSortFilterProxyModel): index = self.sourceModel().index(row, 0, parent) data = index.internalPointer().data(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS) + # virtual address may be empty if not data: return False + # convert virtual address str to int ea = int(data, 16) if self.min_ea <= ea and ea < self.max_ea: @@ -99,7 +112,13 @@ class CapaExplorerRangeProxyModel(QtCore.QSortFilterProxyModel): return False def add_address_range_filter(self, min_ea, max_ea): - """ """ + """add new address range filter + + called when user checks "limit results by current function" in plugin UI + + @param min_ea: minimum virtual address as seen by IDA + @param max_ea: maximum virtual address as seen by IDA + """ self.min_ea = min_ea self.max_ea = max_ea @@ -107,7 +126,10 @@ class CapaExplorerRangeProxyModel(QtCore.QSortFilterProxyModel): self.invalidateFilter() def reset_address_range_filter(self): - """ """ + """remove address range filter (accept all results) + + called when user un-checks "limit results by current function" in plugin UI + """ self.min_ea = None self.max_ea = None self.invalidateFilter() diff --git a/capa/ida/plugin/view.py b/capa/ida/plugin/view.py index 49c3956d..499c6e2d 100644 --- a/capa/ida/plugin/view.py +++ b/capa/ida/plugin/view.py @@ -14,17 +14,16 @@ from capa.ida.plugin.model import CapaExplorerDataModel class CapaExplorerQtreeView(QtWidgets.QTreeView): - """capa explorer QTreeView implementation + """tree view used to display hierarchical capa results - view controls UI action responses and displays data from - CapaExplorerDataModel + view controls UI action responses and displays data from CapaExplorerDataModel - view does not modify CapaExplorerDataModel directly - data - modifications should be implemented in CapaExplorerDataModel + view does not modify CapaExplorerDataModel directly - data modifications should be implemented + in CapaExplorerDataModel """ def __init__(self, model, parent=None): - """ initialize CapaExplorerQTreeView """ + """initialize view""" super(CapaExplorerQtreeView, self).__init__(parent) self.setModel(model) @@ -55,22 +54,21 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): def reset(self): """reset user interface changes - called when view should reset any user interface changes - made since the last reset e.g. IDA window highlighting + called when view should reset any user interface changes made since the last reset e.g. IDA window highlighting """ self.expandToDepth(0) self.resize_columns_to_content() def resize_columns_to_content(self): - """ reset view columns to contents """ + """reset view columns to contents""" self.header().resizeSections(QtWidgets.QHeaderView.ResizeToContents) def map_index_to_source_item(self, model_index): """map proxy model index to source model item - @param model_index: QModelIndex* + @param model_index: QModelIndex - @retval QObject* + @retval QObject """ # assume that self.model here is either: # - CapaExplorerDataModel, or @@ -107,7 +105,7 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): @param data: data passed to slot @param slot: slot to connect - @retval QAction* + @retval QAction """ action = QtWidgets.QAction(display, self.parent) action.setData(data) @@ -120,7 +118,7 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): @param data: tuple - @yield QAction* + @yield QAction """ default_actions = ( ("Copy column", data, self.slot_copy_column), @@ -136,7 +134,7 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): @param data: tuple - @yield QAction* + @yield QAction """ function_actions = (("Rename function", data, self.slot_rename_function),) @@ -153,11 +151,11 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): creates custom context menu containing default actions - @param pos: TODO - @param item: TODO - @param model_index: TODO + @param pos: cursor position + @param item: CapaExplorerDataItem + @param model_index: QModelIndex - @retval QMenu* + @retval QMenu """ menu = QtWidgets.QMenu() @@ -169,14 +167,13 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): 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 + creates custom context menu with both default actions and function actions - @param pos: TODO - @param item: TODO - @param model_index: TODO + @param pos: cursor position + @param item: CapaExplorerDataItem + @param model_index: QModelIndex - @retval QMenu* + @retval QMenu """ menu = QtWidgets.QMenu() @@ -188,8 +185,8 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): def show_custom_context_menu(self, menu, pos): """display custom context menu in view - @param menu: TODO - @param pos: TODO + @param menu: QMenu to display + @param pos: cursor position """ if menu: menu.exec_(self.viewport().mapToGlobal(pos)) @@ -197,10 +194,9 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): def slot_copy_column(self, action): """slot connected to custom context menu - allows user to select a column and copy the data - to clipboard + allows user to select a column and copy the data to clipboard - @param action: QAction* + @param action: QAction """ _, item, model_index = action.data() self.send_data_to_clipboard(item.data(model_index.column())) @@ -208,10 +204,9 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): def slot_copy_row(self, action): """slot connected to custom context menu - allows user to select a row and copy the space-delimited - data to clipboard + allows user to select a row and copy the space-delimited data to clipboard - @param action: QAction* + @param action: QAction """ _, item, _ = action.data() self.send_data_to_clipboard(str(item)) @@ -219,10 +214,9 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): 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 + allows user to select a edit a function name and push changes to IDA - @param action: QAction* + @param action: QAction """ _, item, model_index = action.data() @@ -234,10 +228,9 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): 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 + displays custom context menu to user containing action relevant to the item selected - @param pos: TODO + @param pos: cursor position """ model_index = self.indexAt(pos) item = self.map_index_to_source_item(model_index) @@ -256,9 +249,11 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): self.show_custom_context_menu(menu, pos) def slot_double_click(self, model_index): - """slot connected to double click event + """slot connected to double-click event - @param model_index: QModelIndex* + if address column clicked, navigate IDA to address, else un/expand item clicked + + @param model_index: QModelIndex """ if not model_index.isValid(): return