From 74748ac2c0c1ef7a170192a7bba6f3681a085f55 Mon Sep 17 00:00:00 2001 From: harshit Date: Fri, 29 Nov 2024 23:55:09 +0530 Subject: [PATCH] black fix --- capa/ida/plugin/model.py | 148 +++++++++++++++++++++++++++++++-------- 1 file changed, 117 insertions(+), 31 deletions(-) diff --git a/capa/ida/plugin/model.py b/capa/ida/plugin/model.py index cf2ba4c1b..817f73ae9 100644 --- a/capa/ida/plugin/model.py +++ b/capa/ida/plugin/model.py @@ -53,7 +53,9 @@ def __init__(self, parent=None): """initialize model""" super().__init__(parent) # root node does not have parent, contains header columns - self.root_node = CapaExplorerDataItem(None, ["Rule Information", "Address", "Details"]) + self.root_node = CapaExplorerDataItem( + None, ["Rule Information", "Address", "Details"] + ) def reset(self): """reset UI elements (e.g. checkboxes, IDA color highlights) @@ -62,7 +64,9 @@ def reset(self): """ 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): + for model_index in self.iterateChildrenIndexFromRootIndex( + root_index, ignore_root=False + ): model_index.internalPointer().setChecked(False) self.reset_ida_highlighting(model_index.internalPointer(), False) self.dataChanged.emit(model_index, model_index) @@ -116,7 +120,10 @@ def data(self, model_index, role): # show tooltip containing rule source return item.source - if role == QtCore.Qt.CheckStateRole and column == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION: + if ( + role == QtCore.Qt.CheckStateRole + and column == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION + ): # inform view how to display content of checkbox - un/checked if not item.canCheck(): return None @@ -153,7 +160,10 @@ def data(self, model_index, role): font.setBold(True) return font - if role == QtCore.Qt.ForegroundRole and column == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS: + if ( + role == QtCore.Qt.ForegroundRole + and column == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS + ): # set color for virtual address column return QtGui.QColor(37, 147, 215) @@ -271,7 +281,12 @@ def reset_ida_highlighting(self, item, checked): @param checked: True, item checked, False item not checked """ if not isinstance( - item, (CapaExplorerStringViewItem, CapaExplorerInstructionViewItem, CapaExplorerByteViewItem) + item, + ( + CapaExplorerStringViewItem, + CapaExplorerInstructionViewItem, + CapaExplorerByteViewItem, + ), ): # ignore other item types return @@ -303,10 +318,13 @@ def setData(self, model_index, value, role): if ( role == QtCore.Qt.CheckStateRole - and model_index.column() == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION + and model_index.column() + == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION ): # user un/checked box - un/check parent and children - for child_index in self.iterateChildrenIndexFromRootIndex(model_index, ignore_root=False): + for child_index in self.iterateChildrenIndexFromRootIndex( + model_index, ignore_root=False + ): child_index.internalPointer().setChecked(value) self.reset_ida_highlighting(child_index.internalPointer(), value) self.dataChanged.emit(child_index, child_index) @@ -315,7 +333,8 @@ def setData(self, model_index, value, role): if ( role == QtCore.Qt.EditRole and value - and model_index.column() == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION + and model_index.column() + == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION and isinstance(model_index.internalPointer(), CapaExplorerFunctionItem) ): # user renamed function - update IDA database and data model @@ -371,7 +390,10 @@ def render_capa_doc_statement_node( if statement.description: display += f" ({statement.description})" return CapaExplorerDefaultItem(parent, display) - elif isinstance(statement, rd.CompoundStatement) and statement.type == rd.CompoundStatementType.NOT: + elif ( + isinstance(statement, rd.CompoundStatement) + and statement.type == rd.CompoundStatementType.NOT + ): # TODO(mike-hunhoff): verify that we can display NOT statements # https://github.com/mandiant/capa/issues/1602 pass @@ -403,7 +425,9 @@ def render_capa_doc_statement_node( for location in locations: # for each location render child node for range statement - self.render_capa_doc_feature(parent2, match, statement.child, location, doc) + self.render_capa_doc_feature( + parent2, match, statement.child, location, doc + ) return parent2 elif isinstance(statement, rd.SubscopeStatement): @@ -414,7 +438,9 @@ def render_capa_doc_statement_node( else: raise RuntimeError("unexpected match statement type: " + str(statement)) - def render_capa_doc_match(self, parent: CapaExplorerDataItem, match: rd.Match, doc: rd.ResultDocument): + def render_capa_doc_match( + self, parent: CapaExplorerDataItem, match: rd.Match, doc: rd.ResultDocument + ): """render capa match read from doc @param parent: parent node to which new child is assigned @@ -427,17 +453,28 @@ def render_capa_doc_match(self, parent: CapaExplorerDataItem, match: rd.Match, d return # optional statement with no successful children is empty - if isinstance(match.node, rd.StatementNode) and match.node.statement.type == rd.CompoundStatementType.OPTIONAL: + if ( + isinstance(match.node, rd.StatementNode) + and match.node.statement.type == rd.CompoundStatementType.OPTIONAL + ): if not any(m.success for m in match.children): return if isinstance(match.node, rd.StatementNode): parent2 = self.render_capa_doc_statement_node( - parent, match, match.node.statement, [addr.to_capa() for addr in match.locations], doc + parent, + match, + match.node.statement, + [addr.to_capa() for addr in match.locations], + doc, ) elif isinstance(match.node, rd.FeatureNode): parent2 = self.render_capa_doc_feature_node( - parent, match, match.node.feature, [addr.to_capa() for addr in match.locations], doc + parent, + match, + match.node.feature, + [addr.to_capa() for addr in match.locations], + doc, ) else: raise RuntimeError("unexpected node type: " + str(match.node.type)) @@ -447,7 +484,9 @@ def render_capa_doc_match(self, parent: CapaExplorerDataItem, match: rd.Match, d def render_capa_doc_by_function(self, doc: rd.ResultDocument): """render rule matches by function meaning each rule match is nested under function where it was found""" - matches_by_function: dict[AbsoluteVirtualAddress, tuple[CapaExplorerFunctionItem, set[str]]] = {} + matches_by_function: dict[ + AbsoluteVirtualAddress, tuple[CapaExplorerFunctionItem, set[str]] + ] = {} for rule in rutils.capability_rules(doc): match_eas: list[int] = [] @@ -468,7 +507,9 @@ def render_capa_doc_by_function(self, doc: rd.ResultDocument): # create a new function root to nest its rule matches; Note: we must use the address of the # function here so everything is displayed properly matches_by_function[func_address] = ( - CapaExplorerFunctionItem(self.root_node, func_address, can_check=False), + CapaExplorerFunctionItem( + self.root_node, func_address, can_check=False + ), set(), ) @@ -484,7 +525,13 @@ def render_capa_doc_by_function(self, doc: rd.ResultDocument): func_root, rule.meta.name, rule.meta.namespace or "", - len([ea for ea in match_eas if capa.ida.helpers.get_func_start_ea(ea) == func_ea]), + len( + [ + ea + for ea in match_eas + if capa.ida.helpers.get_func_start_ea(ea) == func_ea + ] + ), rule.source, can_check=False, ) @@ -494,7 +541,13 @@ def render_capa_doc_by_program(self, doc: rd.ResultDocument): for rule in rutils.capability_rules(doc): rule_name = rule.meta.name rule_namespace = rule.meta.namespace or "" - parent = CapaExplorerRuleItem(self.root_node, rule_name, rule_namespace, len(rule.matches), rule.source) + parent = CapaExplorerRuleItem( + self.root_node, + rule_name, + rule_namespace, + len(rule.matches), + rule.source, + ) for location_, match in rule.matches: location = location_.to_capa() @@ -509,7 +562,9 @@ def render_capa_doc_by_program(self, doc: rd.ResultDocument): elif capa.rules.Scope.INSTRUCTION in rule.meta.scopes: parent2 = CapaExplorerInstructionItem(parent, location) else: - raise RuntimeError("unexpected rule scope: " + str(rule.meta.scopes.static)) + raise RuntimeError( + "unexpected rule scope: " + str(rule.meta.scopes.static) + ) self.render_capa_doc_match(parent2, match, doc) @@ -535,14 +590,33 @@ def capa_doc_feature_to_display(self, feature: frzf.Feature) -> str: @param feature: capa feature read from doc """ # Use the specific type from the feature instead of direct string assignment - FeatureType = Union[Literal[ - 'os', 'arch', 'format', 'match', 'characteristic', - 'export', 'import', 'section', 'function name', - 'substring', 'regex', 'string', 'class', 'namespace', - 'api', 'property', 'number', 'bytes', 'offset', - 'mnemonic', 'operand number', 'operand offset', - 'basic block' - ]] + FeatureType = Union[ + Literal[ + "os", + "arch", + "format", + "match", + "characteristic", + "export", + "import", + "section", + "function name", + "substring", + "regex", + "string", + "class", + "namespace", + "api", + "property", + "number", + "bytes", + "offset", + "mnemonic", + "operand number", + "operand offset", + "basic block", + ] + ] key: FeatureType = feature.type value = feature.dict(by_alias=True).get(feature.type) @@ -641,7 +715,9 @@ def render_capa_doc_feature( if matched_rule is not None: matched_rule_source = matched_rule.source - return CapaExplorerRuleMatchItem(parent, display, source=matched_rule_source) + return CapaExplorerRuleMatchItem( + parent, display, source=matched_rule_source + ) elif isinstance(feature, (frzf.RegexFeature, frzf.SubstringFeature)): for capture, addrs in sorted(match.captures.items()): @@ -649,7 +725,10 @@ def render_capa_doc_feature( assert isinstance(addr, frz.Address) if location == addr.value: return CapaExplorerStringViewItem( - parent, display, location, '"' + capa.features.common.escape_string(capture) + '"' + parent, + display, + location, + '"' + capa.features.common.escape_string(capture) + '"', ) # programming error: the given location should always be found in the regex matches @@ -680,7 +759,10 @@ def render_capa_doc_feature( elif isinstance(feature, frzf.StringFeature): # display string preview return CapaExplorerStringViewItem( - parent, display, location, f'"{capa.features.common.escape_string(feature.string)}"' + parent, + display, + location, + f'"{capa.features.common.escape_string(feature.string)}"', ) elif isinstance( @@ -722,7 +804,11 @@ def update_function_name(self, old_name, new_name): # recursive search for all instances of old function name for model_index in self.match( - root_index, QtCore.Qt.DisplayRole, old_name, hits=-1, flags=QtCore.Qt.MatchRecursive + root_index, + QtCore.Qt.DisplayRole, + old_name, + hits=-1, + flags=QtCore.Qt.MatchRecursive, ): if not isinstance(model_index.internalPointer(), CapaExplorerFunctionItem): continue