Reworked format parsing to align better with the rest of capa logic.

capa/features/common.py

@@ -420,7 +420,7 @@ def __init__(self, value: str, description=None):
 FORMAT_SC32 = "sc32"
 FORMAT_SC64 = "sc64"
 FORMAT_FREEZE = "freeze"
-FORMAT_CS = "script_cs"
+FORMAT_SCRIPT = "script"
 FORMAT_UNKNOWN = "unknown"
 
 

capa/features/extractors/script.py

@@ -1,6 +1,7 @@
+import os
 from typing import Tuple, Iterator
 
-from capa.features.common import OS, OS_ANY, ARCH_ANY, FORMAT_CS, Arch, Feature, ScriptLanguage
+from capa.features.common import OS, OS_ANY, ARCH_ANY, Arch, Feature, ScriptLanguage
 from capa.features.address import NO_ADDRESS, Address, FileOffsetRangeAddress
 
 LANG_CS = "c_sharp"
@@ -18,7 +19,8 @@ def extract_os() -> Iterator[Tuple[Feature, Address]]:
     yield OS(OS_ANY), NO_ADDRESS
 
 
-def get_language_from_format(format_: str) -> str:
-    if format_ == FORMAT_CS:
+def get_language_from_ext(path: str):
+    _, ext = os.path.splitext(path)
+    if ext == ".cs":
         return LANG_CS
-    return "unknown"
+    raise ValueError("{path} has an unrecognized or an unsupported extension.")

capa/features/extractors/ts/extractor.py

@@ -8,10 +8,10 @@
 
 
 class TreeSitterFeatureExtractor(FeatureExtractor):
-    def __init__(self, path: str, format_: str):
+    def __init__(self, path: str):
         super().__init__()
         self.path = path
-        self.language = capa.features.extractors.script.get_language_from_format(format_)
+        self.language = capa.features.extractors.script.get_language_from_ext(path)
         with open(self.path, "rb") as f:
             self.buf = f.read()
         self.engine = capa.features.extractors.ts.engine.TreeSitterExtractorEngine(self.language)

capa/helpers.py

@@ -10,16 +10,12 @@
 from typing import NoReturn
 
 from capa.exceptions import UnsupportedFormatError
-from capa.features.common import FORMAT_CS, FORMAT_SC32, FORMAT_SC64, FORMAT_UNKNOWN
+from capa.features.common import FORMAT_SC32, FORMAT_SC64, FORMAT_SCRIPT, FORMAT_UNKNOWN
 
 EXTENSIONS_SHELLCODE_32 = ("sc32", "raw32")
 EXTENSIONS_SHELLCODE_64 = ("sc64", "raw64")
-<<<<<<< HEAD
 EXTENSIONS_ELF = "elf_"
-=======
-EXTENSION_CS = "cs"
-
->>>>>>> Added initial capa control flow for scripts in C#.
+EXTENSIONS_SUPPORTED_SCRIPTS = "cs"
 
 logger = logging.getLogger("capa")
 
@@ -56,8 +52,8 @@ def get_format_from_extension(sample: str) -> str:
         return FORMAT_SC32
     elif sample.endswith(EXTENSIONS_SHELLCODE_64):
         return FORMAT_SC64
-    elif sample.endswith(EXTENSION_CS):
-        return FORMAT_CS
+    elif sample.endswith(EXTENSIONS_SUPPORTED_SCRIPTS):
+        return FORMAT_SCRIPT
     return FORMAT_UNKNOWN
 
 

capa/main.py

@@ -65,6 +65,7 @@
     FORMAT_SC64,
     FORMAT_DOTNET,
     FORMAT_FREEZE,
+    FORMAT_SCRIPT,
 )
 from capa.features.address import NO_ADDRESS
 from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor
@@ -345,11 +346,11 @@ def has_file_limitation(rules: RuleSet, capabilities: MatchResults, is_standalon
     return False
 
 
-def is_supported_script(format_: str):
+def is_script_format(format_: str):
     """
     If the script format was recognized, then it is supported.
     """
-    return format_.startswith("script")
+    return format_ == FORMAT_SCRIPT
 
 
 def is_supported_format(sample: str) -> bool:
@@ -521,10 +522,10 @@ def get_extractor(
       UnsupportedArchError
       UnsupportedOSError
     """
-    if is_supported_script(format_):
+    if format_ == FORMAT_SCRIPT:
         import capa.features.extractors.ts.extractor
 
-        return capa.features.extractors.ts.extractor.TreeSitterFeatureExtractor(path, format_)
+        return capa.features.extractors.ts.extractor.TreeSitterFeatureExtractor(path)
 
     if format_ not in (FORMAT_SC32, FORMAT_SC64):
         if not is_supported_format(path):
@@ -705,7 +706,7 @@ def collect_metadata(
 
     format_ = get_format(sample_path)
 
-    if is_supported_script(format_):
+    if format_ == FORMAT_SCRIPT:
         arch = get_script_arch()
         os_ = get_script_os()
     else: