Fix: Issue #2307 (#2439)

* fix #2307 --------- Co-authored-by: Moritz <[email protected]>
mandiant · Dec 5, 2024 · 28c0234 · 28c0234
1 parent f57f909
commit 28c0234
Show file tree

Hide file tree

Showing 5 changed files with 91 additions and 63 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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:]