Skip to content

Commit

Permalink
type: Fix feature type annotations in capa IDA plugin
Browse files Browse the repository at this point in the history
Resolve type incompatibility by introducing FeatureType type alias with strict literal types
  • Loading branch information
harshit-wadhwani committed Nov 29, 2024
1 parent 20654c2 commit e11313c
Showing 1 changed file with 11 additions and 2 deletions.
13 changes: 11 additions & 2 deletions capa/ida/plugin/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -529,12 +529,21 @@ 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
# 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
value = feature.dict(by_alias=True).get(feature.type)

if value:
Expand Down

0 comments on commit e11313c

Please sign in to comment.