From b77e68df190ef0934d2f5643395a987478e6ca39 Mon Sep 17 00:00:00 2001
From: Yacine Elhamer <elhamer.yacine@gmail.com>
Date: Thu, 22 Jun 2023 14:17:06 +0100
Subject: [PATCH] fix codestyle and typing

---
 capa/helpers.py          |  2 +-
 capa/main.py             | 19 +++++++++++++------
 scripts/show-features.py |  6 ++----
 3 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/capa/helpers.py b/capa/helpers.py
index d06c6676..676e1ceb 100644
--- a/capa/helpers.py
+++ b/capa/helpers.py
@@ -14,7 +14,7 @@ from typing import NoReturn
 import tqdm
 
 from capa.exceptions import UnsupportedFormatError
-from capa.features.common import FORMAT_PE, FORMAT_SC32, FORMAT_SC64, FORMAT_CAPE, FORMAT_DOTNET, FORMAT_UNKNOWN, Format
+from capa.features.common import FORMAT_PE, FORMAT_CAPE, FORMAT_SC32, FORMAT_SC64, FORMAT_DOTNET, FORMAT_UNKNOWN, Format
 
 EXTENSIONS_SHELLCODE_32 = ("sc32", "raw32")
 EXTENSIONS_SHELLCODE_64 = ("sc64", "raw64")
diff --git a/capa/main.py b/capa/main.py
index 0b6372a2..9b3e4bf9 100644
--- a/capa/main.py
+++ b/capa/main.py
@@ -20,7 +20,7 @@ import textwrap
 import itertools
 import contextlib
 import collections
-from typing import Any, Dict, List, Tuple, Callable
+from typing import Any, Dict, List, Tuple, Union, Callable
 
 import halo
 import tqdm
@@ -43,13 +43,13 @@ import capa.render.vverbose
 import capa.features.extractors
 import capa.render.result_document
 import capa.render.result_document as rdoc
-import capa.features.extractors.cape.extractor
 import capa.features.extractors.common
 import capa.features.extractors.pefile
 import capa.features.extractors.dnfile_
 import capa.features.extractors.elffile
 import capa.features.extractors.dotnetfile
 import capa.features.extractors.base_extractor
+import capa.features.extractors.cape.extractor
 from capa.rules import Rule, Scope, RuleSet
 from capa.engine import FeatureSet, MatchResults
 from capa.helpers import (
@@ -70,15 +70,21 @@ from capa.features.common import (
     FORMAT_ELF,
     OS_WINDOWS,
     FORMAT_AUTO,
+    FORMAT_CAPE,
     FORMAT_SC32,
     FORMAT_SC64,
-    FORMAT_CAPE,
     FORMAT_DOTNET,
     FORMAT_FREEZE,
     FORMAT_RESULT,
 )
 from capa.features.address import NO_ADDRESS, Address
-from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor
+from capa.features.extractors.base_extractor import (
+    BBHandle,
+    InsnHandle,
+    FunctionHandle,
+    DynamicExtractor,
+    FeatureExtractor,
+)
 
 RULES_PATH_DEFAULT_STRING = "(embedded rules)"
 SIGNATURES_PATH_DEFAULT_STRING = "(embedded signatures)"
@@ -518,7 +524,7 @@ def get_extractor(
     sigpaths: List[str],
     should_save_workspace=False,
     disable_progress=False,
-) -> FeatureExtractor:
+) -> Union[FeatureExtractor, DynamicExtractor]:
     """
     raises:
       UnsupportedFormatError
@@ -537,9 +543,10 @@ def get_extractor(
             raise UnsupportedOSError()
 
     if format_ == FORMAT_CAPE:
-        import capa.features.extractors.cape.extractor
         import json
 
+        import capa.features.extractors.cape.extractor
+
         with open(path, "r+", encoding="utf-8") as f:
             report = json.load(f)
         return capa.features.extractors.cape.extractor.CapeExtractor.from_report(report)
diff --git a/scripts/show-features.py b/scripts/show-features.py
index 7cc93dda..c8ed2251 100644
--- a/scripts/show-features.py
+++ b/scripts/show-features.py
@@ -135,13 +135,11 @@ def main(argv=None):
             log_unsupported_runtime_error()
             return -1
 
-
     if dynamic:
         dynamic_analysis(extractor, args)
     else:
         static_analysis(extractor, args)
 
-
     return 0
 
 
@@ -188,7 +186,7 @@ def dynamic_analysis(extractor: capa.features.extractors.base_extractor.DynamicE
         if args.process not in [ph.inner["name"] for ph in args.process]:
             print(f"{args.process} not a process")
             return -1
-        
+
     print_process_features(process_handles, extractor)
 
 
@@ -255,6 +253,7 @@ def print_function_features(functions, extractor: capa.features.extractors.base_
                         # may be an issue while piping to less and encountering non-ascii characters
                         continue
 
+
 def ida_main():
     import idc
 
@@ -284,7 +283,6 @@ def ida_main():
     return 0
 
 
-
 if __name__ == "__main__":
     if capa.main.is_runtime_ida():
         ida_main()