From 28c0234339468483f9d21f833ff4a461d5550ef7 Mon Sep 17 00:00:00 2001 From: Harshit Wadhwani <77691984+harshit-wadhwani@users.noreply.github.com> Date: Thu, 5 Dec 2024 14:23:15 +0530 Subject: [PATCH] Fix: Issue #2307 (#2439) * fix #2307 --------- Co-authored-by: Moritz --- capa/features/freeze/features.py | 101 ++++++++++++++++--------------- capa/ida/plugin/model.py | 47 +++++++++++--- capa/render/vverbose.py | 2 +- scripts/compare-backends.py | 1 - scripts/inspect-binexport2.py | 3 - 5 files changed, 91 insertions(+), 63 deletions(-) diff --git a/capa/features/freeze/features.py b/capa/features/freeze/features.py index b3d01f08c..ffff1acd5 100644 --- a/capa/features/freeze/features.py +++ b/capa/features/freeze/features.py @@ -6,7 +6,7 @@ # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and limitations under the License. import binascii -from typing import Union, Optional +from typing import Union, Literal, Optional, Annotated from pydantic import Field, BaseModel, ConfigDict @@ -209,168 +209,171 @@ def feature_from_capa(f: capa.features.common.Feature) -> "Feature": class OSFeature(FeatureModel): - type: str = "os" + type: Literal["os"] = "os" os: str description: Optional[str] = None class ArchFeature(FeatureModel): - type: str = "arch" + type: Literal["arch"] = "arch" arch: str description: Optional[str] = None class FormatFeature(FeatureModel): - type: str = "format" + type: Literal["format"] = "format" format: str description: Optional[str] = None class MatchFeature(FeatureModel): - type: str = "match" + type: Literal["match"] = "match" match: str description: Optional[str] = None class CharacteristicFeature(FeatureModel): - type: str = "characteristic" + type: Literal["characteristic"] = "characteristic" characteristic: str description: Optional[str] = None class ExportFeature(FeatureModel): - type: str = "export" + type: Literal["export"] = "export" export: str description: Optional[str] = None class ImportFeature(FeatureModel): - type: str = "import" + type: Literal["import"] = "import" import_: str = Field(alias="import") description: Optional[str] = None class SectionFeature(FeatureModel): - type: str = "section" + type: Literal["section"] = "section" section: str description: Optional[str] = None class FunctionNameFeature(FeatureModel): - type: str = "function name" + type: Literal["function name"] = "function name" function_name: str = Field(alias="function name") description: Optional[str] = None class SubstringFeature(FeatureModel): - type: str = "substring" + type: Literal["substring"] = "substring" substring: str description: Optional[str] = None class RegexFeature(FeatureModel): - type: str = "regex" + type: Literal["regex"] = "regex" regex: str description: Optional[str] = None class StringFeature(FeatureModel): - type: str = "string" + type: Literal["string"] = "string" string: str description: Optional[str] = None class ClassFeature(FeatureModel): - type: str = "class" + type: Literal["class"] = "class" class_: str = Field(alias="class") description: Optional[str] = None class NamespaceFeature(FeatureModel): - type: str = "namespace" + type: Literal["namespace"] = "namespace" namespace: str description: Optional[str] = None class BasicBlockFeature(FeatureModel): - type: str = "basic block" + type: Literal["basic block"] = "basic block" description: Optional[str] = None class APIFeature(FeatureModel): - type: str = "api" + type: Literal["api"] = "api" api: str description: Optional[str] = None class PropertyFeature(FeatureModel): - type: str = "property" + type: Literal["property"] = "property" access: Optional[str] = None property: str description: Optional[str] = None class NumberFeature(FeatureModel): - type: str = "number" + type: Literal["number"] = "number" number: Union[int, float] description: Optional[str] = None class BytesFeature(FeatureModel): - type: str = "bytes" + type: Literal["bytes"] = "bytes" bytes: str description: Optional[str] = None class OffsetFeature(FeatureModel): - type: str = "offset" + type: Literal["offset"] = "offset" offset: int description: Optional[str] = None class MnemonicFeature(FeatureModel): - type: str = "mnemonic" + type: Literal["mnemonic"] = "mnemonic" mnemonic: str description: Optional[str] = None class OperandNumberFeature(FeatureModel): - type: str = "operand number" + type: Literal["operand number"] = "operand number" index: int operand_number: int = Field(alias="operand number") description: Optional[str] = None class OperandOffsetFeature(FeatureModel): - type: str = "operand offset" + type: Literal["operand offset"] = "operand offset" index: int operand_offset: int = Field(alias="operand offset") description: Optional[str] = None -Feature = Union[ - OSFeature, - ArchFeature, - FormatFeature, - MatchFeature, - CharacteristicFeature, - ExportFeature, - ImportFeature, - SectionFeature, - FunctionNameFeature, - SubstringFeature, - RegexFeature, - StringFeature, - ClassFeature, - NamespaceFeature, - APIFeature, - PropertyFeature, - NumberFeature, - BytesFeature, - OffsetFeature, - MnemonicFeature, - OperandNumberFeature, - OperandOffsetFeature, - # Note! this must be last, see #1161 - BasicBlockFeature, +Feature = Annotated[ + Union[ + OSFeature, + ArchFeature, + FormatFeature, + MatchFeature, + CharacteristicFeature, + ExportFeature, + ImportFeature, + SectionFeature, + FunctionNameFeature, + SubstringFeature, + RegexFeature, + StringFeature, + ClassFeature, + NamespaceFeature, + APIFeature, + PropertyFeature, + NumberFeature, + BytesFeature, + OffsetFeature, + MnemonicFeature, + OperandNumberFeature, + OperandOffsetFeature, + # Note! this must be last, see #1161 + BasicBlockFeature, + ], + Field(discriminator="type"), ] diff --git a/capa/ida/plugin/model.py b/capa/ida/plugin/model.py index 0d8221b12..5d5731318 100644 --- a/capa/ida/plugin/model.py +++ b/capa/ida/plugin/model.py @@ -271,7 +271,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 @@ -433,11 +438,19 @@ def render_capa_doc_match(self, parent: CapaExplorerDataItem, match: rd.Match, d 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)) @@ -494,7 +507,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() @@ -529,12 +548,12 @@ def render_capa_doc(self, doc: rd.ResultDocument, by_function: bool): # inform model changes have ended self.endResetModel() - def capa_doc_feature_to_display(self, feature: frzf.Feature): + def capa_doc_feature_to_display(self, feature: frzf.Feature) -> str: """convert capa doc feature type string to display string for ui @param feature: capa feature read from doc """ - key = feature.type + key = str(feature.type) value = feature.dict(by_alias=True).get(feature.type) if value: @@ -640,7 +659,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 @@ -671,7 +693,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( @@ -713,7 +738,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 diff --git a/capa/render/vverbose.py b/capa/render/vverbose.py index c021bfbb0..1c2255e32 100644 --- a/capa/render/vverbose.py +++ b/capa/render/vverbose.py @@ -168,7 +168,7 @@ def render_feature( ): console.write(" " * indent) - key = feature.type + key = str(feature.type) value: Optional[str] if isinstance(feature, frzf.BasicBlockFeature): # i don't think it makes sense to have standalone basic block features. diff --git a/scripts/compare-backends.py b/scripts/compare-backends.py index fa4ddb010..37336654c 100644 --- a/scripts/compare-backends.py +++ b/scripts/compare-backends.py @@ -125,7 +125,6 @@ def collect(args): key = str(file) for backend in BACKENDS: - if (backend, file.name) in { ("binja", "0953cc3b77ed2974b09e3a00708f88de931d681e2d0cb64afbaf714610beabe6.exe_") }: diff --git a/scripts/inspect-binexport2.py b/scripts/inspect-binexport2.py index 07fc79eca..2b59e35e9 100644 --- a/scripts/inspect-binexport2.py +++ b/scripts/inspect-binexport2.py @@ -75,7 +75,6 @@ def _render_expression_tree( tree_index: int, o: io.StringIO, ): - expression_index = operand.expression_index[tree_index] expression = be2.expression[expression_index] children_tree_indexes: list[int] = expression_tree[tree_index] @@ -124,7 +123,6 @@ def _render_expression_tree( return elif expression.type == BinExport2.Expression.OPERATOR: - if len(children_tree_indexes) == 1: # prefix operator, like "ds:" if expression.symbol != "!": @@ -250,7 +248,6 @@ def inspect_instruction(be2: BinExport2, instruction: BinExport2.Instruction, ad def main(argv=None): - if argv is None: argv = sys.argv[1:]