ida/plugin/form.py: replace usage of '==' with usage of 'in' operator

mandiant · Jul 12, 2023 · 1703039 · 1703039
1 parent 34d3d6c
commit 1703039
Showing 1 changed file with 13 additions and 8 deletions.
diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py
@@ -1192,10 +1192,15 @@ def update_rule_status(self, rule_text: str):
             return
 
         is_match: bool = False
-        if self.rulegen_current_function is not None and rule.scopes in (
-            capa.rules.Scope.FUNCTION,
-            capa.rules.Scope.BASIC_BLOCK,
-            capa.rules.Scope.INSTRUCTION,
+        if self.rulegen_current_function is not None and any(
+            [
+                s in rule.scopes
+                for s in (
+                    capa.rules.Scope.FUNCTION,
+                    capa.rules.Scope.BASIC_BLOCK,
+                    capa.rules.Scope.INSTRUCTION,
+                )
+            ]
         ):
             try:
                 _, func_matches, bb_matches, insn_matches = self.rulegen_feature_cache.find_code_capabilities(
@@ -1205,13 +1210,13 @@ def update_rule_status(self, rule_text: str):
                 self.set_rulegen_status(f"Failed to create function rule matches from rule set ({e})")
                 return
 
-            if rule.scopes == capa.rules.Scope.FUNCTION and rule.name in func_matches.keys():
+            if capa.rules.Scope.FUNCTION in rule.scopes and rule.name in func_matches.keys():
                 is_match = True
-            elif rule.scopes == capa.rules.Scope.BASIC_BLOCK and rule.name in bb_matches.keys():
+            elif capa.rules.Scope.BASIC_BLOCK in rules.scopes and rule.name in bb_matches.keys():
                 is_match = True
-            elif rule.scopes == capa.rules.Scope.INSTRUCTION and rule.name in insn_matches.keys():
+            elif capa.rules.Scope.INSTRUCTION in rules.scopes and rule.name in insn_matches.keys():
                 is_match = True
-        elif rule.scopes == capa.rules.Scope.FILE:
+        elif capa.rules.Scope.FILE in rules.scopes:
             try:
                 _, file_matches = self.rulegen_feature_cache.find_file_capabilities(ruleset)
             except Exception as e: