Skip to content

Commit

Permalink
black fix
Browse files Browse the repository at this point in the history
  • Loading branch information
@harshit-wadhwani
harshit-wadhwani committed Nov 29, 2024
1 parent 9a0c33c commit 74748ac
Showing 1 changed file with 117 additions and 31 deletions.
148 changes: 117 additions & 31 deletions capa/ida/plugin/model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand All @@ -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))
Expand All @@ -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] = []

Expand All @@ -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(),
)

Expand All @@ -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,
)
Expand All @@ -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()
Expand All @@ -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)

Expand All @@ -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)

Expand Down Expand Up @@ -641,15 +715,20 @@ 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()):
for addr in addrs:
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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 74748ac

Please sign in to comment.