use proper str type

mandiant · Dec 4, 2024 · 26363cd · 26363cd
1 parent 4103b10
commit 26363cd
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 31 deletions.
diff --git a/capa/ida/plugin/model.py b/capa/ida/plugin/model.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.
 
-from typing import Union, Literal, Optional
+from typing import Optional
 from collections import deque
 
 import idc
@@ -553,35 +553,7 @@ 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",
-            ]
-        ]
-        key: FeatureType = feature.type
+        key = str(feature.type)
         value = feature.dict(by_alias=True).get(feature.type)
 
         if value:

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.