Added test cases for TS Extractor.

capa/features/extractors/script.py

@@ -24,7 +24,6 @@ def extract_format() -> Iterator[Tuple[Feature, Address]]:
 
 
 def get_language_from_ext(path: str):
-    _, ext = os.path.splitext(path)
-    if ext == ".cs":
+    if path.endswith((".cs", ".cs_")):
         return LANG_CS
     raise ValueError(f"{path} has an unrecognized or an unsupported extension.")

tests/fixtures.py

@@ -13,7 +13,7 @@
 import itertools
 import contextlib
 import collections
-from typing import Set, Dict
+from typing import Set, Dict, Union
 from functools import lru_cache
 
 import pytest
@@ -38,6 +38,7 @@
     Feature,
 )
 from capa.features.address import Address
+from capa.features.extractors.ts.extractor import TreeSitterFeatureExtractor
 from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle
 from capa.features.extractors.dnfile.extractor import DnfileFeatureExtractor
 
@@ -177,6 +178,13 @@ def get_ts_extractor_engine(language, path):
     return capa.features.extractors.ts.engine.TreeSitterExtractorEngine(language, path)
 
 
+@lru_cache(maxsize=1)
+def get_ts_extractor(path):
+    import capa.features.extractors.ts.extractor
+
+    return capa.features.extractors.ts.extractor.TreeSitterFeatureExtractor(path)
+
+
 def extract_global_features(extractor):
     features = collections.defaultdict(set)
     for feature, va in extractor.extract_global_features():
@@ -359,9 +367,13 @@ def sample(request):
     return resolve_sample(request.param)
 
 
-def get_function(extractor, fva: int) -> FunctionHandle:
+def get_function(extractor, fva: Union[int, tuple]) -> FunctionHandle:
+    if isinstance(fva, tuple) and not isinstance(extractor, TreeSitterFeatureExtractor):
+        raise ValueError("invalid fva format")
     for fh in extractor.get_functions():
-        if isinstance(extractor, DnfileFeatureExtractor):
+        if isinstance(extractor, TreeSitterFeatureExtractor):
+            addr = (fh.inner.start_byte, fh.inner.end_byte)
+        elif isinstance(extractor, DnfileFeatureExtractor):
             addr = fh.inner.offset
         else:
             addr = fh.address
@@ -475,6 +487,37 @@ def scope(request):
     return resolve_scope(request.param)
 
 
+def resolve_scope_ts(scope):
+    if scope == "global":
+        inner_fn = lambda extractor: extract_global_features(extractor)
+    elif scope == "file":
+
+        def inner_fn(extractor):
+            features = extract_file_features(extractor)
+            for k, vs in extract_global_features(extractor).items():
+                features[k].update(vs)
+            return features
+
+    elif scope.startswith("function"):
+        # like `function=(155, 192)`
+        def inner_fn(extractor):
+            fh = get_function(extractor, eval(scope.partition("=")[2]))
+            features = extract_function_features(extractor, fh)
+            for k, vs in extract_global_features(extractor).items():
+                features[k].update(vs)
+            return features
+
+    else:
+        raise ValueError("unexpected scope fixture")
+    inner_fn.__name__ = scope
+    return inner_fn
+
+
+@pytest.fixture
+def scope_ts(request):
+    return resolve_scope_ts(request.param)
+
+
 def make_test_id(values):
     return "-".join(map(str, values))
 

tests/test_ts.py

@@ -1,10 +1,13 @@
 from typing import List, Tuple
 
 import pytest
+import fixtures
 from fixtures import *
 from tree_sitter import Node, Tree
 
-from capa.features.address import FileOffsetRangeAddress
+from capa.features.file import Import
+from capa.features.common import OS, OS_ANY, ARCH_ANY, FORMAT_SCRIPT, Arch, Format, String, Namespace, ScriptLanguage
+from capa.features.address import NO_ADDRESS, FileOffsetRangeAddress
 from capa.features.extractors.script import LANG_CS
 from capa.features.extractors.ts.query import QueryBinding
 from capa.features.extractors.ts.engine import TreeSitterExtractorEngine
@@ -26,10 +29,6 @@ def do_test_range(engine: TreeSitterExtractorEngine, node: Node, expected_range:
     assert engine.get_range(node).startswith(expected_range) if startswith else engine.get_range(node) == expected_range
 
 
-def do_test_id_range(engine: TreeSitterExtractorEngine, node: Node, expected_id_range: str, startswith: bool = False):
-    do_test_range(engine, engine.get_object_id(node), expected_id_range, startswith)
-
-
 def do_test_range_address(engine: TreeSitterExtractorEngine, node: Node):
     assert isinstance(engine.get_address(node), FileOffsetRangeAddress)
     addr = engine.get_address(node)
@@ -235,8 +234,6 @@ def do_test_ts_engine_function_names_parsing(
                 "global statements": [
                     'string stdout = "";',
                     'string stderr = "";',
-                    "void die() {",
-                    "void Page_Load(object sender, System.EventArgs e) {",
                 ],
                 "all import names": ["System.Diagnostics.ProcessStartInfo", "System.Diagnostics.Process"],
                 "all function names": [],
@@ -259,3 +256,24 @@ def test_ts_engine(request: pytest.FixtureRequest, engine_str: str, expected_dic
     do_test_ts_engine_global_statements_parsing(engine, expected_dict["global statements"])
     do_test_ts_engine_namespaces_parsing(engine, expected_dict["namespaces"])
     do_test_ts_engine_default_range_address(engine)
+
+
+FEATURE_PRESENCE_TESTS_SCRIPTS = sorted(
+    [
+        ("cs_f397cb", "global", Arch(ARCH_ANY), True),
+        ("cs_f397cb", "global", OS(OS_ANY), True),
+        ("cs_f397cb", "file", Format(FORMAT_SCRIPT), True),
+        ("cs_f397cb", "file", ScriptLanguage(LANG_CS), True),
+        ("cs_f397cb", "file", Namespace("System"), True),
+        ("cs_f397cb", "file", String(""), True),
+        ("cs_f397cb", "function=(0x38,0x16c)", String("Not Found"), True),
+        ("cs_f397cb", "function=(0x16e,0x7ce)", String("127.0.0.1"), True),
+        ("cs_f397cb", "function=(0x16e,0x7ce)", Import("System.Diagnostics.ProcessStartInfo"), True),
+        ("cs_f397cb", "function=(0x16e,0x7ce)", Import("System.Diagnostics.Process"), True),
+    ]
+)
+
+
+@parametrize("sample, scope_ts, feature, expected", FEATURE_PRESENCE_TESTS_SCRIPTS, indirect=["sample", "scope_ts"])
+def test_ts_extractor(sample, scope_ts, feature, expected):
+    fixtures.do_test_feature_presence(fixtures.get_ts_extractor, sample, scope_ts, feature, expected)