Bug fixes and code refactoring of the Tree Sitter extractor.

capa/features/extractors/ts/build.py

@@ -1,4 +1,10 @@
+from tree_sitter import Language
+
 build_dir = "build/my-languages.so"
 languages = [
     "vendor/tree-sitter-c-sharp",
 ]
+
+
+def ts_build():
+    Language.build_library(build_dir, languages)

capa/features/extractors/ts/engine.py

@@ -1,8 +1,9 @@
-from typing import List, Tuple
+from typing import List, Tuple, Iterator
 
 from tree_sitter import Node, Tree, Parser
 
 import capa.features.extractors.ts.sig
+import capa.features.extractors.ts.build
 import capa.features.extractors.ts.query
 from capa.features.address import FileOffsetRangeAddress
 from capa.features.extractors.ts.query import QueryBinding
@@ -16,6 +17,7 @@ class TreeSitterExtractorEngine:
     import_signatures: set
 
     def __init__(self, language: str, path: str):
+        capa.features.extractors.ts.build.ts_build()
         self.language = language
         self.query = capa.features.extractors.ts.query.QueryBinding(language)
         self.import_signatures = capa.features.extractors.ts.sig.load_import_signatures(language)
@@ -27,6 +29,9 @@ def __init__(self, language: str, path: str):
     def get_language(self):
         return self.language
 
+    def get_ts_language(self):
+        return self.query.language
+
     def parse(self):
         self.parser = Parser()
         self.parser.set_language(self.get_ts_language())
@@ -38,37 +43,46 @@ def get_new_objects(self, node: Node) -> List[Tuple[Node, str]]:
     def get_object_id(self, node: Node) -> Node:
         return node.child_by_field_name(self.query.new_object_field_name)
 
+    def get_new_object_ids(self, node: Node) -> Iterator[Node]:
+        for obj_node, _ in self.get_new_objects(node):
+            yield self.get_object_id(obj_node)
+
     def get_import_names(self, node: Node) -> List[Tuple[Node, str]]:
         join_names = capa.features.extractors.ts.sig.get_name_joiner(self.language)
         import_names = []
         namespaces = set([self.get_range(node) for node, _ in self.get_all_namespaces()])
-        for node, _ in self.get_new_objects(node):
+        for node, _ in self.get_new_object_ids(node):
             for namespace in namespaces:
                 name = join_names(namespace, self.get_range(node))
                 if name in self.import_signatures:
                     import_names.append(name)
         return import_names
 
-    def get_functions(self, node: Node) -> List[Tuple[Node, str]]:
-        return self.query.function_def.captures(node)
-
-    def get_all_functions(self) -> List[Tuple[Node, str]]:
-        return self.get_functions(self.tree.root_node)
+    def get_function_definitions(self, node: Node = None) -> List[Tuple[Node, str]]:
+        return self.query.function_definition.captures(node if node is not None else self.tree.root_node)
 
     def get_function_definition_id(self, node: Node) -> Node:
-        return node.child_by_field_name(self.query.function_def_field_name)
+        return node.child_by_field_name(self.query.function_definition_field_name)
+
+    def get_function_definition_ids(self, node: Node) -> Iterator[Node]:
+        for fn_node, _ in self.get_function_definitions(node):
+            yield self.get_function_definition_id(fn_node)
 
     def get_function_calls(self, node: Node) -> List[Tuple[Node, str]]:
         return self.query.function_call.captures(node)
 
     def get_function_call_id(self, node: Node) -> Node:
         return node.child_by_field_name(self.query.function_call_field_name)
 
+    def get_function_call_ids(self, node: Node) -> Iterator[Node]:
+        for fn_node, _ in self.get_function_calls(node):
+            yield self.get_function_call_id(fn_node)
+
     def get_function_names(self, node: Node) -> List[Tuple[Node, str]]:
         join_names = capa.features.extractors.ts.sig.get_name_joiner(self.language)
         function_names = []
         namespaces = set([self.get_range(node) for node, _ in self.get_all_namespaces()])
-        for node, _ in self.get_function_calls(node):
+        for node, _ in self.get_function_call_ids(node):
             for namespace in namespaces:
                 name = join_names(namespace, self.get_range(node))
                 if name in self.import_signatures:
@@ -97,4 +111,4 @@ def get_address(self, node: Node):
         return FileOffsetRangeAddress(node.start_byte, node.end_byte)
 
     def get_default_address(self):
-        return self.get_addr(self.tree.root_node)
+        return self.get_address(self.tree.root_node)

capa/features/extractors/ts/extractor.py

@@ -27,7 +27,7 @@ def extract_file_features(self) -> Iterator[Tuple[Feature, Address]]:
         yield from capa.features.extractors.ts.file.extract_features(self.engine)
 
     def get_functions(self) -> Iterator[FunctionHandle]:
-        for node, _ in self.engine.get_all_functions():
+        for node, _ in self.engine.get_function_definitions():
             yield FunctionHandle(address=self.engine.get_address(node), inner=node)
 
     def extract_function_features(self, f: FunctionHandle) -> Iterator[Tuple[Feature, Address]]:

capa/features/extractors/ts/query.py

@@ -12,8 +12,8 @@ class QueryBinding:
     language: Language
     new_object: Query
     new_object_field_name: str
-    function_def: Query
-    function_def_field_name: str
+    function_definition: Query
+    function_definition_field_name: str
     function_call: Query
     function_call_field_name: str
     string_literal: Query
@@ -26,8 +26,8 @@ def __init__(self, language: str):
         if language == LANG_CS:
             self.new_object = self.language.query("(object_creation_expression) @object.new")
             self.new_object_field_name = "type"
-            self.function_def = self.language.query("(local_function_statement) @function.def")
-            self.function_def_field_name = "name"
+            self.function_definition = self.language.query("(local_function_statement) @function.definition")
+            self.function_definition_field_name = "name"
             self.function_call = self.language.query("(invocation_expression) @function.call")
             self.function_call_field_name = "function"
             self.string_literal = self.language.query("(string_literal) @string-literal")