From 725361c94908b212014679f193983433fca63d43 Mon Sep 17 00:00:00 2001 From: Michael Hunhoff Date: Mon, 14 Sep 2020 09:41:57 -0600 Subject: [PATCH 1/5] add progress indicator wait box --- capa/ida/plugin/form.py | 261 +++++++++++++++++++++++++++++++--------- 1 file changed, 203 insertions(+), 58 deletions(-) diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index d5fd82d4..d7e294fc 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -12,6 +12,7 @@ import logging import collections import idaapi +import ida_kernwin import ida_settings from PyQt5 import QtGui, QtCore, QtWidgets @@ -30,6 +31,43 @@ logger = logging.getLogger(__name__) settings = ida_settings.IDASettings("capa") +class UserCancelledError(Exception): + pass + + +class CapaExplorerProgressIndicator(QtCore.QObject): + progress = QtCore.pyqtSignal(str) + + def __init__(self): + """ """ + super(CapaExplorerProgressIndicator, self).__init__() + + def update(self, text): + """ """ + if ida_kernwin.user_cancelled(): + raise UserCancelledError("user cancelled") + self.progress.emit("extracting features from %s" % text) + + +class CapaExplorerFeatureExtractor(capa.features.extractors.ida.IdaFeatureExtractor): + def __init__(self, progress): + super(CapaExplorerFeatureExtractor, self).__init__() + self.progress = progress + + def extract_function_features(self, f): + self.progress.update("function at 0x%x" % f.start_ea) + for (feature, ea) in capa.features.extractors.ida.function.extract_features(f): + yield feature, ea + + def extract_basic_block_features(self, f, bb): + for (feature, ea) in capa.features.extractors.ida.basicblock.extract_features(f, bb): + yield feature, ea + + def extract_insn_features(self, f, bb, insn): + for (feature, ea) in capa.features.extractors.ida.insn.extract_features(f, bb, insn): + yield feature, ea + + class CapaExplorerForm(idaapi.PluginForm): """form element for plugin interface""" @@ -361,74 +399,176 @@ class CapaExplorerForm(idaapi.PluginForm): """run capa analysis and render results in UI""" # new analysis, new doc self.doc = None + self.process_total = 0 + self.process_count = 0 - # resolve rules directory - check self and settings first, then ask user - if not self.rule_path: - if "rule_path" in settings and os.path.exists(settings["rule_path"]): - self.rule_path = settings["rule_path"] - else: - rule_path = self.ask_user_directory() - if not rule_path: - capa.ida.helpers.inform_user_ida_ui( - "You must select a file directory containing capa rules to start analysis" - ) - logger.warning( - "No rules loaded, cannot start analysis. You can download the standard collection of capa rules from https://github.com/fireeye/capa-rules." - ) - self.set_view_status_label("No rules loaded.") - self.disable_controls() - return - self.rule_path = rule_path - settings.user["rule_path"] = rule_path + self.set_view_status_label("No rules loaded") + self.disable_controls() + + def update_wait_box(text): + """ """ + ida_kernwin.replace_wait_box("Processing; %s" % text) + + def slot_progress_feature_extraction(text): + """ """ + update_wait_box("%s (%d/%d)" % (text, self.process_count, self.process_total)) + self.process_count += 1 + + progress = CapaExplorerProgressIndicator() + progress.progress.connect(slot_progress_feature_extraction) + extractor = CapaExplorerFeatureExtractor(progress) + + update_wait_box("calculating analysis") try: - rules = capa.main.get_rules(self.rule_path, True) - rule_count = len(rules) - 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 capa rules from %s" % self.rule_path) - logger.error( - "Failed to load rules from %s (%s). Make sure your file directory contains properly formatted capa rules. You can download the standard collection of capa rules from https://github.com/fireeye/capa-rules.", - self.rule_path, - e, - ) - self.rule_path = "" - settings.user.del_value("rule_path") - self.set_view_status_label("No rules loaded") - self.disable_controls() + self.process_total += len(tuple(extractor.get_functions())) + except Exception as e: + logger.error("Failed to calculate analysis (Error: %s)." % e) return - meta = capa.ida.helpers.collect_metadata() + if ida_kernwin.user_cancelled(): + logger.info("User cancelled analysis.") + return - capabilities, counts = capa.main.find_capabilities( - rules, capa.features.extractors.ida.IdaFeatureExtractor(), True - ) - meta["analysis"].update(counts) + update_wait_box("loading rules") - # support binary files specifically for x86/AMD64 shellcode - # 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( - " capa currently only supports analyzing binary files containing x86/AMD64 shellcode with IDA." + try: + # resolve rules directory - check self and settings first, then ask user + if not self.rule_path: + if "rule_path" in settings and os.path.exists(settings["rule_path"]): + self.rule_path = settings["rule_path"] + else: + idaapi.info("You must select a file directory containing capa rules before running analysis.") + rule_path = self.ask_user_directory() + if not rule_path: + logger.warning( + "You must select a file directory containing capa rules before running analysis. The standard collection of capa rules can be downloaded from https://github.com/fireeye/capa-rules." + ) + return + self.rule_path = rule_path + settings.user["rule_path"] = rule_path + except Exception as e: + logger.error("Failed to load capa rules (Error: %s)." % e) + return + + if ida_kernwin.user_cancelled(): + logger.info("User cancelled analysis.") + return + + rule_path = self.rule_path + + try: + if not os.path.exists(rule_path): + raise IOError("rule path %s does not exist or cannot be accessed" % rule_path) + + rule_paths = [] + if os.path.isfile(rule_path): + rule_paths.append(rule_path) + elif os.path.isdir(rule_path): + for root, dirs, files in os.walk(rule_path): + if ".github" in root: + # the .github directory contains CI config in capa-rules + # this includes some .yml files + # these are not rules + continue + for file in files: + if not file.endswith(".yml"): + if not (file.endswith(".md") or file.endswith(".git") or file.endswith(".txt")): + # expect to see readme.md, format.md, and maybe a .git directory + # other things maybe are rules, but are mis-named. + logger.warning("skipping non-.yml file: %s", file) + continue + rule_path = os.path.join(root, file) + rule_paths.append(rule_path) + + rules = [] + total_paths = len(rule_paths) + for (i, rule_path) in enumerate(rule_paths): + update_wait_box("loading rule %d/%d from %s" % (i + 1, total_paths, self.rule_path)) + if ida_kernwin.user_cancelled(): + raise UserCancelledError("user cancelled") + try: + rule = capa.rules.Rule.from_yaml_file(rule_path) + except capa.rules.InvalidRule: + raise + else: + rule.meta["capa/path"] = rule_path + if capa.main.is_nursery_rule_path(rule_path): + rule.meta["capa/nursery"] = True + rules.append(rule) + + rule_count = len(rules) + rules = capa.rules.RuleSet(rules) + except UserCancelledError: + logger.info("User cancelled analysis.") + return + except Exception as e: + capa.ida.helpers.inform_user_ida_ui("Failed to load capa rules from %s" % self.rule_path) + logger.error("Failed to load rules from %s (Error: %s).", self.rule_path, e) + logger.error("Make sure your file directory contains properly formatted capa rules. You can download the standard collection of capa rules from https://github.com/fireeye/capa-rules.") + self.rule_path = "" + settings.user.del_value("rule_path") + return + + if ida_kernwin.user_cancelled(): + logger.info("User cancelled analysis.") + return + + ida_kernwin.replace_wait_box("Processing; extracting features") + + try: + meta = capa.ida.helpers.collect_metadata() + capabilities, counts = capa.main.find_capabilities( + rules, extractor, disable_progress=True ) - 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) + meta["analysis"].update(counts) + except UserCancelledError: + logger.info("User cancelled analysis.") + return + except Exception as e: + logger.error("Failed to extract capabilities from database (Error: %s)" % e) + return - capa.ida.helpers.inform_user_ida_ui("capa encountered file type warnings during analysis") + ida_kernwin.replace_wait_box("Processing; checking for file limitations") - if capa.main.has_file_limitation(rules, capabilities, is_standalone=False): - capa.ida.helpers.inform_user_ida_ui("capa encountered file limitation warnings during analysis") + try: + # support binary files specifically for x86/AMD64 shellcode + # 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( + " capa currently only supports analyzing binary files containing x86/AMD64 shellcode with IDA." + ) + 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.doc = capa.render.convert_capabilities_to_result_document(meta, rules, capabilities) - self.model_data.render_capa_doc(self.doc) - self.render_capa_doc_mitre_summary() + capa.ida.helpers.inform_user_ida_ui("capa encountered file type warnings during analysis") + + if capa.main.has_file_limitation(rules, capabilities, is_standalone=False): + capa.ida.helpers.inform_user_ida_ui("capa encountered file limitation warnings during analysis") + except Exception as e: + logger.error("Failed to check for file limitations (Error: %s)" % e) + return + + if ida_kernwin.user_cancelled(): + logger.info("User cancelled analysis.") + return + + update_wait_box("Processing; rendering results") + + try: + self.doc = capa.render.convert_capabilities_to_result_document(meta, rules, capabilities) + self.model_data.render_capa_doc(self.doc) + self.render_capa_doc_mitre_summary() + except Exception as e: + logger.error("Failed to render results (Error: %s)" % e) + return self.enable_controls() self.set_view_status_label("Loaded %d capa rules from %s" % (rule_count, self.rule_path)) @@ -510,8 +650,13 @@ class CapaExplorerForm(idaapi.PluginForm): self.range_model_proxy.invalidate() self.search_model_proxy.invalidate() self.model_data.clear() + + ida_kernwin.show_wait_box("Processing") self.load_capa_results() + ida_kernwin.hide_wait_box() + self.ida_reset() + logger.info("Analysis completed.") def slot_reset(self, checked): @@ -562,7 +707,7 @@ class CapaExplorerForm(idaapi.PluginForm): """create Qt dialog to ask user for a directory""" return str( QtWidgets.QFileDialog.getExistingDirectory( - self.parent, "Please select a file directory containing capa rules", self.rule_path + self.parent, "Please select a capa rules directory", self.rule_path ) ) From 0d93df7d59d9ecc68e2b4f8545d2846e9ee69946 Mon Sep 17 00:00:00 2001 From: Michael Hunhoff Date: Mon, 14 Sep 2020 11:29:17 -0600 Subject: [PATCH 2/5] updating documentation --- capa/ida/plugin/form.py | 42 ++++++++++++++++++++++++++++------------- 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index d7e294fc..23c66fd6 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -32,24 +32,33 @@ settings = ida_settings.IDASettings("capa") class UserCancelledError(Exception): + """throw exception when user cancels action""" pass class CapaExplorerProgressIndicator(QtCore.QObject): + """implement progress signal, used during feature extraction""" progress = QtCore.pyqtSignal(str) def __init__(self): - """ """ + """initialize signal object""" super(CapaExplorerProgressIndicator, self).__init__() def update(self, text): - """ """ + """emit progress update + + check if user cancelled action, raise exception for parent function to catch + """ if ida_kernwin.user_cancelled(): raise UserCancelledError("user cancelled") self.progress.emit("extracting features from %s" % text) class CapaExplorerFeatureExtractor(capa.features.extractors.ida.IdaFeatureExtractor): + """subclass the IdaFeatureExtractor + + track progress during feature extraction, also allow user to cancel feature extraction + """ def __init__(self, progress): super(CapaExplorerFeatureExtractor, self).__init__() self.progress = progress @@ -77,6 +86,8 @@ class CapaExplorerForm(idaapi.PluginForm): self.form_title = name self.rule_path = "" + self.process_total = 0 + self.process_count = 0 self.parent = None self.ida_hooks = None @@ -396,21 +407,26 @@ class CapaExplorerForm(idaapi.PluginForm): self.slot_analyze() def load_capa_results(self): - """run capa analysis and render results in UI""" + """run capa analysis and render results in UI + + note: this function must always return, exception or not, in order for plugin to safely close the IDA + wait box + """ # new analysis, new doc self.doc = None self.process_total = 0 self.process_count = 0 + # default from view self.set_view_status_label("No rules loaded") self.disable_controls() def update_wait_box(text): - """ """ + """update the IDA wait box""" ida_kernwin.replace_wait_box("Processing; %s" % text) def slot_progress_feature_extraction(text): - """ """ + """slot function to handle feature extraction progress updates""" update_wait_box("%s (%d/%d)" % (text, self.process_count, self.process_total)) self.process_count += 1 @@ -423,7 +439,7 @@ class CapaExplorerForm(idaapi.PluginForm): try: self.process_total += len(tuple(extractor.get_functions())) except Exception as e: - logger.error("Failed to calculate analysis (Error: %s)." % e) + logger.error("Failed to calculate analysis (error: %s)." % e) return if ida_kernwin.user_cancelled(): @@ -448,7 +464,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.rule_path = rule_path settings.user["rule_path"] = rule_path except Exception as e: - logger.error("Failed to load capa rules (Error: %s)." % e) + logger.error("Failed to load capa rules (error: %s)." % e) return if ida_kernwin.user_cancelled(): @@ -504,7 +520,7 @@ class CapaExplorerForm(idaapi.PluginForm): return except Exception as e: capa.ida.helpers.inform_user_ida_ui("Failed to load capa rules from %s" % self.rule_path) - logger.error("Failed to load rules from %s (Error: %s).", self.rule_path, e) + logger.error("Failed to load rules from %s (error: %s).", self.rule_path, e) logger.error("Make sure your file directory contains properly formatted capa rules. You can download the standard collection of capa rules from https://github.com/fireeye/capa-rules.") self.rule_path = "" settings.user.del_value("rule_path") @@ -514,7 +530,7 @@ class CapaExplorerForm(idaapi.PluginForm): logger.info("User cancelled analysis.") return - ida_kernwin.replace_wait_box("Processing; extracting features") + update_wait_box("extracting features") try: meta = capa.ida.helpers.collect_metadata() @@ -526,10 +542,10 @@ class CapaExplorerForm(idaapi.PluginForm): logger.info("User cancelled analysis.") return except Exception as e: - logger.error("Failed to extract capabilities from database (Error: %s)" % e) + logger.error("Failed to extract capabilities from database (error: %s)" % e) return - ida_kernwin.replace_wait_box("Processing; checking for file limitations") + update_wait_box("checking for file limitations") try: # support binary files specifically for x86/AMD64 shellcode @@ -553,7 +569,7 @@ class CapaExplorerForm(idaapi.PluginForm): if capa.main.has_file_limitation(rules, capabilities, is_standalone=False): capa.ida.helpers.inform_user_ida_ui("capa encountered file limitation warnings during analysis") except Exception as e: - logger.error("Failed to check for file limitations (Error: %s)" % e) + logger.error("Failed to check for file limitations (error: %s)" % e) return if ida_kernwin.user_cancelled(): @@ -567,7 +583,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.model_data.render_capa_doc(self.doc) self.render_capa_doc_mitre_summary() except Exception as e: - logger.error("Failed to render results (Error: %s)" % e) + logger.error("Failed to render results (error: %s)" % e) return self.enable_controls() From 89e409157fb7a46908806d2faf39907122dcd2d2 Mon Sep 17 00:00:00 2001 From: Michael Hunhoff Date: Mon, 14 Sep 2020 13:00:06 -0600 Subject: [PATCH 3/5] updating progress message --- capa/ida/plugin/form.py | 103 ++++++++++++++++++++-------------------- 1 file changed, 51 insertions(+), 52 deletions(-) diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index 23c66fd6..d5338234 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -33,11 +33,13 @@ settings = ida_settings.IDASettings("capa") class UserCancelledError(Exception): """throw exception when user cancels action""" + pass class CapaExplorerProgressIndicator(QtCore.QObject): """implement progress signal, used during feature extraction""" + progress = QtCore.pyqtSignal(str) def __init__(self): @@ -59,23 +61,16 @@ class CapaExplorerFeatureExtractor(capa.features.extractors.ida.IdaFeatureExtrac track progress during feature extraction, also allow user to cancel feature extraction """ - def __init__(self, progress): + + def __init__(self): super(CapaExplorerFeatureExtractor, self).__init__() - self.progress = progress + self.indicator = CapaExplorerProgressIndicator() def extract_function_features(self, f): - self.progress.update("function at 0x%x" % f.start_ea) + self.indicator.update("function at 0x%X" % f.start_ea) for (feature, ea) in capa.features.extractors.ida.function.extract_features(f): yield feature, ea - def extract_basic_block_features(self, f, bb): - for (feature, ea) in capa.features.extractors.ida.basicblock.extract_features(f, bb): - yield feature, ea - - def extract_insn_features(self, f, bb, insn): - for (feature, ea) in capa.features.extractors.ida.insn.extract_features(f, bb, insn): - yield feature, ea - class CapaExplorerForm(idaapi.PluginForm): """form element for plugin interface""" @@ -221,7 +216,7 @@ class CapaExplorerForm(idaapi.PluginForm): """load status label""" label = QtWidgets.QLabel() label.setAlignment(QtCore.Qt.AlignLeft) - label.setText("Analyze database to get started...") + label.setText("Click Analyze to get started...") self.view_status_label = label @@ -417,34 +412,29 @@ class CapaExplorerForm(idaapi.PluginForm): self.process_total = 0 self.process_count = 0 - # default from view - self.set_view_status_label("No rules loaded") - self.disable_controls() - def update_wait_box(text): """update the IDA wait box""" - ida_kernwin.replace_wait_box("Processing; %s" % text) + ida_kernwin.replace_wait_box("capa explorer...%s" % text) def slot_progress_feature_extraction(text): """slot function to handle feature extraction progress updates""" - update_wait_box("%s (%d/%d)" % (text, self.process_count, self.process_total)) + update_wait_box("%s (%d of %d)" % (text, self.process_count, self.process_total)) self.process_count += 1 - progress = CapaExplorerProgressIndicator() - progress.progress.connect(slot_progress_feature_extraction) - extractor = CapaExplorerFeatureExtractor(progress) + extractor = CapaExplorerFeatureExtractor() + extractor.indicator.progress.connect(slot_progress_feature_extraction) update_wait_box("calculating analysis") try: self.process_total += len(tuple(extractor.get_functions())) except Exception as e: - logger.error("Failed to calculate analysis (error: %s)." % e) - return + logger.error("Failed to calculate analysis (error: %s).", e) + return False if ida_kernwin.user_cancelled(): logger.info("User cancelled analysis.") - return + return False update_wait_box("loading rules") @@ -454,22 +444,22 @@ class CapaExplorerForm(idaapi.PluginForm): if "rule_path" in settings and os.path.exists(settings["rule_path"]): self.rule_path = settings["rule_path"] else: - idaapi.info("You must select a file directory containing capa rules before running analysis.") + idaapi.info("Please select a file directory containing capa rules.") rule_path = self.ask_user_directory() if not rule_path: logger.warning( - "You must select a file directory containing capa rules before running analysis. The standard collection of capa rules can be downloaded from https://github.com/fireeye/capa-rules." + "You must select a file directory containing capa rules before analysis can be run. The standard collection of capa rules can be downloaded from https://github.com/fireeye/capa-rules." ) - return + return False self.rule_path = rule_path settings.user["rule_path"] = rule_path except Exception as e: - logger.error("Failed to load capa rules (error: %s)." % e) - return + logger.error("Failed to load capa rules (error: %s).", e) + return False if ida_kernwin.user_cancelled(): logger.info("User cancelled analysis.") - return + return False rule_path = self.rule_path @@ -500,7 +490,7 @@ class CapaExplorerForm(idaapi.PluginForm): rules = [] total_paths = len(rule_paths) for (i, rule_path) in enumerate(rule_paths): - update_wait_box("loading rule %d/%d from %s" % (i + 1, total_paths, self.rule_path)) + update_wait_box("loading capa rules from %s (%d of %d)" % (self.rule_path, i + 1, total_paths)) if ida_kernwin.user_cancelled(): raise UserCancelledError("user cancelled") try: @@ -517,33 +507,33 @@ class CapaExplorerForm(idaapi.PluginForm): rules = capa.rules.RuleSet(rules) except UserCancelledError: logger.info("User cancelled analysis.") - return + return False except Exception as e: capa.ida.helpers.inform_user_ida_ui("Failed to load capa rules from %s" % self.rule_path) logger.error("Failed to load rules from %s (error: %s).", self.rule_path, e) - logger.error("Make sure your file directory contains properly formatted capa rules. You can download the standard collection of capa rules from https://github.com/fireeye/capa-rules.") + logger.error( + "Make sure your file directory contains properly formatted capa rules. You can download the standard collection of capa rules from https://github.com/fireeye/capa-rules." + ) self.rule_path = "" settings.user.del_value("rule_path") - return + return False if ida_kernwin.user_cancelled(): logger.info("User cancelled analysis.") - return + return False update_wait_box("extracting features") try: meta = capa.ida.helpers.collect_metadata() - capabilities, counts = capa.main.find_capabilities( - rules, extractor, disable_progress=True - ) + capabilities, counts = capa.main.find_capabilities(rules, extractor, disable_progress=True) meta["analysis"].update(counts) except UserCancelledError: logger.info("User cancelled analysis.") - return + return False except Exception as e: - logger.error("Failed to extract capabilities from database (error: %s)" % e) - return + logger.error("Failed to extract capabilities from database (error: %s)", e) + return False update_wait_box("checking for file limitations") @@ -561,7 +551,9 @@ class CapaExplorerForm(idaapi.PluginForm): 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( + " If you don't know the input file type, you can try using the `file` utility to guess it." + ) logger.warning("-" * 80) capa.ida.helpers.inform_user_ida_ui("capa encountered file type warnings during analysis") @@ -569,25 +561,26 @@ class CapaExplorerForm(idaapi.PluginForm): if capa.main.has_file_limitation(rules, capabilities, is_standalone=False): capa.ida.helpers.inform_user_ida_ui("capa encountered file limitation warnings during analysis") except Exception as e: - logger.error("Failed to check for file limitations (error: %s)" % e) - return + logger.error("Failed to check for file limitations (error: %s)", e) + return False if ida_kernwin.user_cancelled(): logger.info("User cancelled analysis.") - return + return False - update_wait_box("Processing; rendering results") + update_wait_box("rendering results") try: self.doc = capa.render.convert_capabilities_to_result_document(meta, rules, capabilities) self.model_data.render_capa_doc(self.doc) self.render_capa_doc_mitre_summary() + self.enable_controls() + self.set_view_status_label("capa rules directory: %s (%d rules)" % (self.rule_path, rule_count)) except Exception as e: - logger.error("Failed to render results (error: %s)" % e) - return + logger.error("Failed to render results (error: %s)", e) + return False - self.enable_controls() - self.set_view_status_label("Loaded %d capa rules from %s" % (rule_count, self.rule_path)) + return True def render_capa_doc_mitre_summary(self): """render MITRE ATT&CK results""" @@ -667,10 +660,16 @@ class CapaExplorerForm(idaapi.PluginForm): self.search_model_proxy.invalidate() self.model_data.clear() - ida_kernwin.show_wait_box("Processing") - self.load_capa_results() + self.disable_controls() + self.set_view_status_label("Loading...") + + ida_kernwin.show_wait_box("capa explorer") + success = self.load_capa_results() ida_kernwin.hide_wait_box() + if not success: + self.set_view_status_label("Click Analyze to get started...") + self.ida_reset() logger.info("Analysis completed.") From a45dbba4b1880d62d3aabb1dbce08519ae267a58 Mon Sep 17 00:00:00 2001 From: Michael Hunhoff Date: Mon, 14 Sep 2020 14:30:27 -0600 Subject: [PATCH 4/5] bug fixes for program rebase hook --- capa/ida/plugin/form.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index d5338234..d9ccad60 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -135,7 +135,7 @@ class CapaExplorerForm(idaapi.PluginForm): ensure any plugin modifications (e.g. hooks and UI changes) are reset before the plugin is closed """ self.unload_ida_hooks() - self.ida_reset() + self.model_data.reset() def load_interface(self): """load user interface""" @@ -398,8 +398,12 @@ class CapaExplorerForm(idaapi.PluginForm): @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 program rebase") - self.slot_analyze() + if idaapi.get_imagebase() != meta.get("prev_base", -1): + capa.ida.helpers.inform_user_ida_ui("Running capa analysis again after program rebase") + self.slot_analyze() + else: + meta["prev_base"] = idaapi.get_imagebase() + self.model_data.reset() def load_capa_results(self): """run capa analysis and render results in UI @@ -641,14 +645,13 @@ class CapaExplorerForm(idaapi.PluginForm): item.setFont(font) return item - def ida_reset(self): - """reset plugin UI + def reset_view_tree(self): + """reset tree view UI controls called when user selects plugin reset from menu """ self.view_limit_results_by_function.setChecked(False) self.view_search_bar.setText("") - self.model_data.reset() self.view_tree.reset_ui() def slot_analyze(self): @@ -658,8 +661,8 @@ class CapaExplorerForm(idaapi.PluginForm): """ self.range_model_proxy.invalidate() self.search_model_proxy.invalidate() + self.model_data.reset() self.model_data.clear() - self.disable_controls() self.set_view_status_label("Loading...") @@ -667,19 +670,21 @@ class CapaExplorerForm(idaapi.PluginForm): success = self.load_capa_results() ida_kernwin.hide_wait_box() + self.reset_view_tree() + if not success: self.set_view_status_label("Click Analyze to get started...") - - self.ida_reset() - - logger.info("Analysis completed.") + logger.info("Analysis failed.") + else: + logger.info("Analysis completed.") def slot_reset(self, checked): """reset UI elements e.g. checkboxes and IDA highlighting """ - self.ida_reset() + self.model_data.reset() + self.reset_view_tree() logger.info("Reset completed.") def slot_checkbox_limit_by_changed(self, state): From a02235e89422f74e235ee1b9c0e45bcd4a58618b Mon Sep 17 00:00:00 2001 From: Michael Hunhoff Date: Mon, 14 Sep 2020 15:12:35 -0600 Subject: [PATCH 5/5] PR change requests --- capa/ida/plugin/form.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index d9ccad60..047dbe91 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -68,8 +68,7 @@ class CapaExplorerFeatureExtractor(capa.features.extractors.ida.IdaFeatureExtrac def extract_function_features(self, f): self.indicator.update("function at 0x%X" % f.start_ea) - for (feature, ea) in capa.features.extractors.ida.function.extract_features(f): - yield feature, ea + return super(CapaExplorerFeatureExtractor, self).extract_function_features(f) class CapaExplorerForm(idaapi.PluginForm):