Implemented the necessary modifications to support embedded templates…

…/html: aspx.

capa/features/extractors/script.py

@@ -26,7 +26,7 @@ def extract_format() -> Iterator[Tuple[Feature, Address]]:
     yield Format(FORMAT_SCRIPT), NO_ADDRESS
 
 
-def get_language_from_ext(path: str):
+def get_language_from_ext(path: str) -> str:
     if path.endswith((".aspx", "aspx_")):
         return LANG_TEM
     if path.endswith((".cs", ".cs_")):

capa/features/extractors/ts/engine.py

@@ -1,14 +1,12 @@
 import re
-from typing import Dict, List, Tuple, Union, Iterator
-from collections import defaultdict
-from dataclasses import dataclass
+from typing import List, Tuple, Iterator, Optional
 
 from tree_sitter import Node, Tree, Parser
 
 import capa.features.extractors.ts.sig
 import capa.features.extractors.ts.build
 from capa.features.address import FileOffsetRangeAddress
-from capa.features.extractors.script import LANG_CS, LANG_JS
+from capa.features.extractors.script import LANG_CS, LANG_JS, LANG_TEM, LANG_HTML
 from capa.features.extractors.ts.query import (
     QueryBinding,
     HTMLQueryBinding,
@@ -21,18 +19,14 @@
 class TreeSitterBaseEngine:
     buf: bytes
     language: str
-    path: str
     query: QueryBinding
     tree: Tree
 
-    def __init__(self, language: str, path: str):
+    def __init__(self, language: str, buf: bytes):
         capa.features.extractors.ts.build.ts_build()
         self.language = language
         self.query = QueryBindingFactory.from_language(language)
-        self.import_signatures = capa.features.extractors.ts.sig.load_import_signatures(language)
-        self.path = path
-        with open(self.path, "rb") as f:
-            self.buf = f.read()
+        self.buf = buf
         self.tree = self.parse()
 
     def parse(self) -> Tree:
@@ -46,19 +40,27 @@ def get_byte_range(self, node: Node) -> bytes:
     def get_range(self, node: Node) -> str:
         return self.get_byte_range(node).decode()
 
-    def get_address(self, node: Node):
+    def get_address(self, node: Node) -> FileOffsetRangeAddress:
         return FileOffsetRangeAddress(node.start_byte, node.end_byte)
 
-    def get_default_address(self):
+    def get_default_address(self) -> FileOffsetRangeAddress:
         return self.get_address(self.tree.root_node)
 
 
 class TreeSitterExtractorEngine(TreeSitterBaseEngine):
     query: ScriptQueryBinding
     import_signatures: set
+    buf_offset: int
+    namespaces: set[str]
+
+    def __init__(self, language: str, buf: bytes, buf_offset: int = 0, additional_namespaces: set[str] = None):
+        super().__init__(language, buf)
+        self.buf_offset = buf_offset
+        self.import_signatures = capa.features.extractors.ts.sig.load_import_signatures(language)
+        self.namespaces = additional_namespaces if additional_namespaces is not None else set()
 
-    def __init__(self, language: str, path: str):
-        super().__init__(language, path)
+    def get_address(self, node: Node) -> FileOffsetRangeAddress:
+        return FileOffsetRangeAddress(self.buf_offset + node.start_byte, self.buf_offset + node.end_byte)
 
     def get_new_objects(self, node: Node) -> List[Tuple[Node, str]]:
         return self.query.new_object.captures(node)
@@ -73,13 +75,13 @@ def get_new_object_ids(self, node: Node) -> Iterator[Node]:
     # TODO: move this elsewhere, does not fit this class
     def get_import_names(self, node: Node) -> Iterator[Tuple[Node, str]]:
         join_names = capa.features.extractors.ts.sig.get_name_joiner(self.language)
-        namespaces = set([self.get_range(ns_node) for ns_node, _ in self.get_namespaces()])
+        self.namespaces = self.namespaces.union(set([self.get_range(ns_node) for ns_node, _ in self.get_namespaces()]))
         for obj_node in self.get_new_object_ids(node):
             obj_name = self.get_range(obj_node)
             if obj_name in self.import_signatures:
                 yield (obj_node, obj_name)
                 continue
-            for namespace in namespaces:
+            for namespace in self.namespaces:
                 obj_name = join_names(namespace, obj_name)
                 if obj_name in self.import_signatures:
                     yield (obj_node, obj_name)
@@ -107,13 +109,13 @@ def get_function_call_ids(self, node: Node) -> Iterator[Node]:
     # TODO: move this elsewhere, does not fit this class
     def get_function_names(self, node: Node) -> Iterator[Tuple[Node, str]]:
         join_names = capa.features.extractors.ts.sig.get_name_joiner(self.language)
-        namespaces = set([self.get_range(ns_node) for ns_node, _ in self.get_namespaces()])
+        self.namespaces = self.namespaces.union(set([self.get_range(ns_node) for ns_node, _ in self.get_namespaces()]))
         for fn_node in self.get_function_call_ids(node):
             fn_name = self.get_range(fn_node)
             if fn_name in self.import_signatures:
                 yield (fn_node, fn_name)
                 continue
-            for namespace in namespaces:
+            for namespace in self.namespaces:
                 fn_name = join_names(namespace, fn_name)
                 if fn_name in self.import_signatures:
                     yield (fn_node, fn_name)
@@ -131,65 +133,73 @@ def get_global_statements(self) -> List[Tuple[Node, str]]:
         return self.query.global_statement.captures(self.tree.root_node)
 
 
-@dataclass
-class ASPXPseudoNode:
-    start_byte: int
-    end_byte: int
-
-
 class TreeSitterTemplateEngine(TreeSitterBaseEngine):
     query: TemplateQueryBinding
 
-    def __init__(self, language: str, path: str):
-        super().__init__(language, path)
+    def __init__(self, buf: bytes):
+        super().__init__(LANG_TEM, buf)
 
     def get_code_sections(self) -> List[Tuple[Node, str]]:
         return self.query.code.captures(self.tree.root_node)
 
+    def get_parsed_code_sections(self) -> Iterator[TreeSitterExtractorEngine]:
+        template_namespaces = set(name for _, name in self.get_template_namespaces())
+        for node, _ in self.get_code_sections():
+            yield TreeSitterExtractorEngine(
+                self.identify_language(), self.get_byte_range(node), node.start_byte, template_namespaces
+            )
+
     def get_content_sections(self) -> List[Tuple[Node, str]]:
         return self.query.content.captures(self.tree.root_node)
 
-    def get_template_namespaces(self) -> Iterator[ASPXPseudoNode]:
+    def identify_language(self) -> str:
+        for node, _ in self.get_code_sections():
+            if self.is_c_sharp(node):
+                return LANG_CS
+        return LANG_JS
+
+    def get_template_namespaces(self) -> Iterator[Tuple[Node, str]]:
         for node, _ in self.get_code_sections():
             if self.is_aspx_import_directive:
-                ns = self.get_aspx_namespace(node)
-                if ns is not None:
-                    yield ns
+                namespace = self.get_aspx_namespace(node)
+                if namespace is not None:
+                    yield node, namespace
 
-    def is_aspx(self, node: Node) -> bool:
-        return self.get_byte_range(node).startswith(b"@")
+    def is_c_sharp(self, node: Node) -> bool:
+        return len(re.findall(r'@ .*Page Language\s*=\s*"C#".*'.encode(), self.get_byte_range(node))) > 0
 
     def is_aspx_import_directive(self, node: Node) -> bool:
         return self.get_byte_range(node).startswith(b"@ Import namespace=")
 
-    def get_aspx_namespace(self, node: Node) -> Union[ASPXPseudoNode, None]:
+    def get_aspx_namespace(self, node: Node) -> Optional[str]:
         match = re.search(r'@ Import namespace="(.*?)"'.encode(), self.get_byte_range(node))
-        if match is None:
-            return None
-        return ASPXPseudoNode(node.start_byte + match.span()[0], node.start_byte + match.span()[1])
+        return match.group().decode() if match is not None else None
 
 
 class TreeSitterHTMLEngine(TreeSitterBaseEngine):
     query: HTMLQueryBinding
+    namespaces: set[str]
 
-    def __init__(self, language: str, path: str):
-        super().__init__(language, path)
+    def __init__(self, buf: bytes, additional_namespaces: set[str] = None):
+        super().__init__(LANG_HTML, buf)
+        self.namespaces = additional_namespaces if additional_namespaces is not None else set()
 
     def get_scripts(self) -> List[Tuple[Node, str]]:
         return self.query.script_element.captures(self.tree.root_node)
 
     def get_attributes(self, node: Node) -> List[Tuple[Node, str]]:
         return self.query.attribute.captures(self.tree.root_node)
 
-    def get_code_sections_by_language(self) -> Dict[str, List[Node]]:
-        code_sections = defaultdict(list)
+    def get_code_sections(self) -> Iterator[Node]:
         for script_node, _ in self.get_scripts():
             for attribute_node, _ in self.get_attributes(script_node):
-                script_language = self.identify_script_language(attribute_node)
-                code_sections[script_language].append(attribute_node)
-        return code_sections
+                yield attribute_node
+
+    def get_parsed_code_sections(self) -> Iterator[TreeSitterExtractorEngine]:
+        for node in self.get_code_sections():
+            yield TreeSitterExtractorEngine(self.identify_language(node), self.get_byte_range(node), node.start_byte)
 
-    def identify_script_language(self, node: Node) -> str:
+    def identify_language(self, node: Node) -> str:
         if self.is_server_side_c_sharp(node):
             return LANG_CS
         return LANG_JS

capa/features/extractors/ts/extractor.py

@@ -1,37 +1,78 @@
-from typing import Tuple, Union, Iterator
+from typing import List, Tuple, Union, Iterator
+
+from tree_sitter import Node
 
 import capa.features.extractors.script
 import capa.features.extractors.ts.file
 import capa.features.extractors.ts.engine
 import capa.features.extractors.ts.global_
 import capa.features.extractors.ts.function
-from capa.features.address import NO_ADDRESS, Address, AbsoluteVirtualAddress
-from capa.features.extractors.ts.engine import TreeSitterExtractorEngine
+from capa.features.common import Namespace
+from capa.features.address import NO_ADDRESS, Address, AbsoluteVirtualAddress, FileOffsetRangeAddress
+from capa.features.extractors.script import LANG_TEM, LANG_HTML
+from capa.features.extractors.ts.engine import TreeSitterHTMLEngine, TreeSitterTemplateEngine, TreeSitterExtractorEngine
 from capa.features.extractors.base_extractor import Feature, BBHandle, InsnHandle, FunctionHandle, FeatureExtractor
 
 
 class TreeSitterFeatureExtractor(FeatureExtractor):
-    engine: TreeSitterExtractorEngine
+    code_sections: List[TreeSitterExtractorEngine]
+    template_namespaces: set[Tuple[Node, str]]
+    language: str
 
     def __init__(self, path: str):
         super().__init__()
-        self.engine = TreeSitterExtractorEngine(capa.features.extractors.script.get_language_from_ext(path), path)
+        self.path = path
+        with open(self.path, "rb") as f:
+            buf = f.read()
+
+        self.language = capa.features.extractors.script.get_language_from_ext(path)
+        if self.language == LANG_TEM:
+            self.code_sections, self.template_namespaces = self.extract_code_from_template(buf)
+        elif self.language == LANG_HTML:
+            self.code_sections = list(self.extract_code_from_html(buf))
+        else:
+            self.code_sections = [TreeSitterExtractorEngine(self.language, buf)]
+
+    def extract_code_from_template(self, buf: bytes) -> Tuple[List[TreeSitterExtractorEngine], set[Tuple[Node, str]]]:
+        template_engine = TreeSitterTemplateEngine(buf)
+        template_namespaces = set(template_engine.get_template_namespaces())
+        code_sections = list(template_engine.get_parsed_code_sections())
+
+        additional_namespaces = set(name for _, name in template_namespaces)
+        for section in template_engine.get_content_sections():
+            section_buf = template_engine.get_byte_range(section)
+            code_sections.extend(list(self.extract_code_from_html(section_buf, additional_namespaces)))
+        return code_sections, template_namespaces
+
+    def extract_code_from_html(
+        self, buf: bytes, additional_namespaces: set[str] = None
+    ) -> Iterator[TreeSitterExtractorEngine]:
+        yield from TreeSitterHTMLEngine(buf, additional_namespaces).get_parsed_code_sections()
 
     def get_base_address(self) -> Union[AbsoluteVirtualAddress, capa.features.address._NoAddress]:
         return NO_ADDRESS
 
+    def extract_template_namespaces(self) -> Iterator[Tuple[Feature, Address]]:
+        for node, name in self.template_namespaces:
+            yield Namespace(name), FileOffsetRangeAddress(node.start_byte, node.end_byte)
+
     def extract_global_features(self) -> Iterator[Tuple[Feature, Address]]:
         yield from capa.features.extractors.ts.global_.extract_features()
 
     def extract_file_features(self) -> Iterator[Tuple[Feature, Address]]:
-        yield from capa.features.extractors.ts.file.extract_features(self.engine)
+        if self.language == LANG_TEM:
+            yield from self.extract_template_namespaces()
+        for engine in self.code_sections:
+            yield from capa.features.extractors.ts.file.extract_features(engine)
 
     def get_functions(self) -> Iterator[FunctionHandle]:
-        for node, _ in self.engine.get_function_definitions():
-            yield FunctionHandle(address=self.engine.get_address(node), inner=node)
+        for engine in self.code_sections:
+            for node, _ in engine.get_function_definitions():
+                yield FunctionHandle(address=engine.get_address(node), inner=node)
 
     def extract_function_features(self, f: FunctionHandle) -> Iterator[Tuple[Feature, Address]]:
-        yield from capa.features.extractors.ts.function.extract_features(f, self.engine)
+        for engine in self.code_sections:
+            yield from capa.features.extractors.ts.function.extract_features(f, engine)
 
     def get_basic_blocks(self, f: FunctionHandle) -> Iterator[BBHandle]:
         yield from []

capa/features/extractors/ts/file.py

@@ -8,10 +8,6 @@
 from capa.features.extractors.ts.engine import TreeSitterExtractorEngine
 
 
-def extract_file_format(engine: TreeSitterExtractorEngine) -> Iterator[Tuple[Feature, Address]]:
-    yield from capa.features.extractors.script.extract_format()
-
-
 def extract_language(engine: TreeSitterExtractorEngine) -> Iterator[Tuple[Feature, Address]]:
     yield from capa.features.extractors.script.extract_language(engine.language, engine.get_default_address())
 
@@ -52,7 +48,6 @@ def extract_features(engine: TreeSitterExtractorEngine) -> Iterator[Tuple[Featur
 
 
 FILE_HANDLERS = (
-    extract_file_format,
     extract_file_function_names,
     extract_file_import_names,
     extract_file_integer_literals,

capa/features/extractors/ts/global_.py

@@ -19,4 +19,8 @@ def extract_features() -> Iterator[Tuple[Feature, Address]]:
             yield feature, addr
 
 
-GLOBAL_HANDLERS = (extract_arch, extract_os)
+def extract_file_format() -> Iterator[Tuple[Feature, Address]]:
+    yield from capa.features.extractors.script.extract_format()
+
+
+GLOBAL_HANDLERS = (extract_arch, extract_os, extract_file_format)

tests/fixtures.py

@@ -172,10 +172,10 @@ def get_dnfile_extractor(path):
 
 
 @lru_cache(maxsize=1)
-def get_ts_extractor_engine(language, path):
+def get_ts_extractor_engine(language, buf):
     import capa.features.extractors.ts.engine
 
-    return capa.features.extractors.ts.engine.TreeSitterExtractorEngine(language, path)
+    return capa.features.extractors.ts.engine.TreeSitterExtractorEngine(language, buf)
 
 
 @lru_cache(maxsize=1)
@@ -963,4 +963,6 @@ def _692f_dotnetfile_extractor():
 
 @pytest.fixture
 def cs_f397cb_extractor_engine():
-    return get_ts_extractor_engine("c_sharp", get_data_path_by_name("cs_f397cb"))
+    with open(get_data_path_by_name("cs_f397cb"), "rb") as f:
+        buf = f.read()
+    return get_ts_extractor_engine("c_sharp", buf)

tests/test_ts.py

@@ -17,7 +17,6 @@ def do_test_ts_engine_init(engine: TreeSitterExtractorEngine):
     assert engine.language == LANG_CS
     assert isinstance(engine.query, QueryBinding)
     assert isinstance(engine.import_signatures, set) and len(engine.import_signatures) > 0
-    assert isinstance(engine.path, str) and len(engine.path) > 0
     assert isinstance(engine.buf, bytes) and len(engine.buf) > 0
     assert isinstance(engine.tree, Tree)
     assert isinstance(engine.get_default_address(), FileOffsetRangeAddress)