Add Bash support to script analysis

Author: saniyafatima07

capa/features/extractors/ts/engine.py

@@ -121,7 +121,7 @@ def get_global_statements(self) -> Iterator[Node]:
 
 
 class TreeSitterExtractorEngine(TreeSitterBaseEngine):
-    query: ScriptQueryBinding
+    query: ScriptQueryBinding | BashQueryBinding
     language_toolkit: LanguageToolkit
     namespaces: set[BaseNamespace]
 
@@ -142,34 +142,47 @@ def __init__(
         self.namespaces = self.namespaces.union(additional_namespaces)
 
     def get_new_object_names(self, node: Node) -> Iterator[Node]:
+        if not isinstance(self.query, ScriptQueryBinding):
+            return
         cursor = QueryCursor(self.query.new_object_name)
         yield from self.get_captured_nodes(cursor, node)
 
     def get_property_names(self, node: Node) -> Iterator[Node]:
+        if not isinstance(self.query, ScriptQueryBinding):
+            return
         cursor = QueryCursor(self.query.property_name)
         yield from self.get_captured_nodes(cursor, node)
 
     def get_processed_property_names(self, node: Node) -> Iterator[Tuple[Node, str]]:
         """Generates captured property name nodes and their associated proper names (see process_property
         for details), e.g.: [(node0, "StartInfo"), (node1, "RedirectStandardOutput")]."""
+        if not isinstance(self.query, ScriptQueryBinding):
+            return
         for pt_node in self.get_property_names(node):
             pt_name = self.language_toolkit.process_property(pt_node, self.get_str(pt_node))
             if pt_name:
                 yield pt_node, pt_name
 
     def get_imported_constants(self, node: Node) -> Iterator[Node]:
+        if not isinstance(self.query, ScriptQueryBinding):
+            return
         cursor = QueryCursor(self.query.imported_constant_name)
         yield from self.get_captured_nodes(cursor, node)
 
     def get_processed_imported_constants(self, node: Node) -> Iterator[Tuple[Node, str]]:
         """Generates captured imported constant nodes and their associated proper names (see process_imported_constant
         for details), e.g.: [(node0, "ssl.CERT_NONE"), (node1, "win32con.FILE_ATTRIBUTE_HIDDEN")]."""
+        if not isinstance(self.query, ScriptQueryBinding):
+            return
         for ic_node in self.get_imported_constants(node):
             ic_name = self.language_toolkit.process_imported_constant(ic_node, self.get_str(ic_node))
             if ic_name:
                 yield ic_node, ic_name
 
     def get_namespaces(self, node: Optional[Node] = None) -> List[Tuple[Node, str]]:
+        if not isinstance(self.query, ScriptQueryBinding):
+            return []
+
         target_node = self.tree.root_node if node is None else node
         cursor = QueryCursor(self.query.namespace)
         namespace_captures: List[Tuple[Node, str]] = []
@@ -181,10 +194,15 @@ def get_namespaces(self, node: Optional[Node] = None) -> List[Tuple[Node, str]]:
         return sorted(namespace_captures, key=self.get_node_capture_sort_key)
 
     def get_processed_namespaces(self, node: Optional[Node] = None) -> Iterator[BaseNamespace]:
+        if not isinstance(self.query, ScriptQueryBinding):
+            return
         for ns_node, query_name in self.get_namespaces(node):
             yield from self.language_toolkit.process_namespace(ns_node, query_name, self.get_str)
 
     def get_direct_method_call(self, node: Node) -> Optional[Node]:
+        if not isinstance(self.query, ScriptQueryBinding):
+            return None
+
         cursor = QueryCursor(self.query.direct_method_call)
         captures = cursor.captures(node)
         for nodes in captures.values():

capa/features/extractors/ts/function.py

@@ -21,6 +21,7 @@
 from capa.features.insn import API, Number, Property
 from capa.features.common import Class, String, Feature, Namespace
 from capa.features.address import Address
+from capa.features.extractors.ts.query import BashQueryBinding
 from capa.features.extractors.ts.tools import BaseNamespace
 from capa.features.extractors.ts.engine import TreeSitterBaseEngine, TreeSitterBashEngine, TreeSitterExtractorEngine
 from capa.features.extractors.base_extractor import FunctionHandle
@@ -167,10 +168,18 @@ def _extract_function_calls(
 ) -> Iterator[Tuple[Feature, Address]]:
     for node in engine.get_function_call_names(fn_node):
         yield from _extract_static_methods(node, engine)
+
+        if isinstance(engine.query, BashQueryBinding):
+            continue
+
         yield from _extract_instance_methods(node, classes, engine)
 
 
 def extract_imports(fn_node: Node, engine: TreeSitterExtractorEngine) -> Iterator[Tuple[Feature, Address]]:
+    if isinstance(engine.query, BashQueryBinding):
+        yield from _extract_function_calls(fn_node, set(), engine)
+        return
+
     classes = {engine.language_toolkit.create_namespace(cls) for cls in get_classes(fn_node, engine)}
     yield from _extract_classes(fn_node, engine)
     yield from _extract_constants(fn_node, engine)

tests/fixtures/__init__.py

@@ -1223,7 +1223,6 @@ def get_ts_template_engine(path):
     return capa.features.extractors.ts.engine.TreeSitterTemplateEngine(buf)
 
 
-@functools.lru_cache(maxsize=1)
 def get_ts_extractor(path):
     import capa.features.extractors.ts.extractor