add support for analysis of source code/scripted languages

.github/mypy/mypy.ini

@@ -76,4 +76,7 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 
 [mypy-dncil.*]
+ignore_missing_imports = True
+
+[mypy-tree_sitter.*]
 ignore_missing_imports = True

capa/features/address.py

@@ -53,6 +53,26 @@ def __repr__(self):
         return f"file(0x{self:x})"
 
 
+class FileOffsetRangeAddress(Address):
+    """an address range relative to the start of a file"""
+
+    def __init__(self, start_byte, end_byte):
+        self.start_byte = start_byte
+        self.end_byte = end_byte
+
+    def __eq__(self, other):
+        return (self.start_byte, self.end_byte) == (self.start_byte, other.end_byte)
+
+    def __lt__(self, other):
+        return (self.start_byte, self.end_byte) < (other.start_byte, other.end_byte)
+
+    def __hash__(self):
+        return hash((self.start_byte, self.end_byte))
+
+    def __repr__(self):
+        return f"file(0x{self.start_byte:x}, 0x{self.end_byte:x})"
+
+
 class DNTokenAddress(Address):
     """a .NET token"""
 

capa/features/common.py

@@ -405,6 +405,12 @@ def __init__(self, value: str, description=None):
         self.name = "os"
 
 
+class ScriptLanguage(Feature):
+    def __init__(self, value: str, description=None):
+        super().__init__(value, description=description)
+        self.name = "script language"
+
+
 FORMAT_PE = "pe"
 FORMAT_ELF = "elf"
 FORMAT_DOTNET = "dotnet"
@@ -414,6 +420,7 @@ def __init__(self, value: str, description=None):
 FORMAT_SC32 = "sc32"
 FORMAT_SC64 = "sc64"
 FORMAT_FREEZE = "freeze"
+FORMAT_SCRIPT = "script"
 FORMAT_UNKNOWN = "unknown"
 
 

capa/features/extractors/script.py

@@ -0,0 +1,40 @@
+import os
+from typing import Tuple, Iterator
+
+from capa.features.common import OS, OS_ANY, ARCH_ANY, FORMAT_SCRIPT, Arch, Format, Feature, ScriptLanguage
+from capa.features.address import NO_ADDRESS, Address, FileOffsetRangeAddress
+
+LANG_CS = "c_sharp"
+LANG_HTML = "html"
+LANG_JS = "javascript"
+LANG_TEM = "embedded_template"
+
+EXT_ASPX = ("aspx", "aspx_")
+EXT_CS = ("cs", "cs_")
+EXT_HTML = ("html", "html_")
+
+
+def extract_arch() -> Iterator[Tuple[Feature, Address]]:
+    yield Arch(ARCH_ANY), NO_ADDRESS
+
+
+def extract_language(language: str, addr: FileOffsetRangeAddress) -> Iterator[Tuple[Feature, Address]]:
+    yield ScriptLanguage(language), addr
+
+
+def extract_os() -> Iterator[Tuple[Feature, Address]]:
+    yield OS(OS_ANY), NO_ADDRESS
+
+
+def extract_format() -> Iterator[Tuple[Feature, Address]]:
+    yield Format(FORMAT_SCRIPT), NO_ADDRESS
+
+
+def get_language_from_ext(path: str) -> str:
+    if path.endswith(EXT_ASPX):
+        return LANG_TEM
+    if path.endswith(EXT_CS):
+        return LANG_CS
+    if path.endswith(EXT_HTML):
+        return LANG_HTML
+    raise ValueError(f"{path} has an unrecognized or an unsupported extension.")

capa/features/extractors/ts/build.py

@@ -0,0 +1,13 @@
+from tree_sitter import Language
+
+build_dir = "build/my-languages.so"
+languages = [
+    "vendor/tree-sitter-c-sharp",
+    "vendor/tree-sitter-embedded-template",
+    "vendor/tree-sitter-html",
+    "vendor/tree-sitter-javascript",
+]
+
+
+def ts_build():
+    Language.build_library(build_dir, languages)

capa/features/extractors/ts/engine.py

@@ -0,0 +1,214 @@
+import re
+from typing import List, Tuple, Iterator, Optional
+
+from tree_sitter import Node, Tree, Parser
+
+import capa.features.extractors.ts.build
+from capa.features.address import FileOffsetRangeAddress
+from capa.features.extractors.script import LANG_CS, LANG_JS, LANG_TEM, LANG_HTML
+from capa.features.extractors.ts.query import (
+    BINDINGS,
+    QueryBinding,
+    HTMLQueryBinding,
+    ScriptQueryBinding,
+    TemplateQueryBinding,
+)
+from capa.features.extractors.ts.tools import LANGUAGE_TOOLKITS, LanguageToolkit
+
+
+class TreeSitterBaseEngine:
+    buf: bytes
+    language: str
+    query: QueryBinding
+    tree: Tree
+
+    def __init__(self, language: str, buf: bytes):
+        capa.features.extractors.ts.build.ts_build()
+        self.language = language
+        self.query = BINDINGS[language]
+        self.buf = buf
+        self.tree = self.parse()
+
+    def parse(self) -> Tree:
+        parser = Parser()
+        parser.set_language(self.query.language)
+        return parser.parse(self.buf)
+
+    def get_byte_range(self, node: Node) -> bytes:
+        return self.buf[node.start_byte : node.end_byte]
+
+    def get_range(self, node: Node) -> str:
+        return self.get_byte_range(node).decode()
+
+    def get_address(self, node: Node) -> FileOffsetRangeAddress:
+        return FileOffsetRangeAddress(node.start_byte, node.end_byte)
+
+    def get_default_address(self) -> FileOffsetRangeAddress:
+        return self.get_address(self.tree.root_node)
+
+
+class TreeSitterExtractorEngine(TreeSitterBaseEngine):
+    query: ScriptQueryBinding
+    language_toolkit: LanguageToolkit
+    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.language_toolkit = LANGUAGE_TOOLKITS[language]
+        self.namespaces = set(self.get_range(ns_node) for ns_node, _ in self.get_namespaces())
+        if additional_namespaces:
+            self.namespaces = self.namespaces.union(additional_namespaces)
+
+    def get_address(self, node: Node) -> FileOffsetRangeAddress:
+        return FileOffsetRangeAddress(self.buf_offset + node.start_byte, self.buf_offset + node.end_byte)
+
+    def get_new_object_names(self, node: Node) -> List[Tuple[Node, str]]:
+        return self.query.new_object_name.captures(node)
+
+    def get_assigned_property_names(self, node: Node) -> List[Tuple[Node, str]]:
+        return self.query.assigned_property_name.captures(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_name(self, node: Node) -> Node:
+        return node.child_by_field_name(self.query.function_definition_field_name)
+
+    def get_function_definition_names(self, node: Node) -> Iterator[Node]:
+        for fn_node, _ in self.get_function_definitions(node):
+            yield self.get_function_definition_name(fn_node)
+
+    def get_function_call_names(self, node: Node) -> List[Tuple[Node, str]]:
+        return self.query.function_call_name.captures(node)
+
+    def get_string_literals(self, node: Node) -> List[Tuple[Node, str]]:
+        return self.query.string_literal.captures(node)
+
+    def get_integer_literals(self, node: Node) -> List[Tuple[Node, str]]:
+        return self.query.integer_literal.captures(node)
+
+    def get_namespaces(self, node: Node = None) -> List[Tuple[Node, str]]:
+        return self.query.namespace.captures(node if node is not None else self.tree.root_node)
+
+    def get_global_statements(self) -> List[Tuple[Node, str]]:
+        return self.query.global_statement.captures(self.tree.root_node)
+
+
+class TreeSitterTemplateEngine(TreeSitterBaseEngine):
+    query: TemplateQueryBinding
+    language_toolkit: LanguageToolkit
+    embedded_language: str
+
+    def __init__(self, buf: bytes):
+        super().__init__(LANG_TEM, buf)
+        self.embedded_language = self.identify_language()
+        self.language_toolkit = LANGUAGE_TOOLKITS[self.embedded_language]
+        self.template_namespaces = set(name for _, name in self.get_template_namespaces())
+
+    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]:
+        for node, _ in self.get_code_sections():
+            # TODO: support JS
+            if self.embedded_language == LANG_CS:
+                yield TreeSitterExtractorEngine(
+                    self.embedded_language,
+                    self.get_byte_range(node),
+                    node.start_byte,
+                    self.template_namespaces,
+                )
+
+    def get_content_sections(self) -> List[Tuple[Node, str]]:
+        return self.query.content.captures(self.tree.root_node)
+
+    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_imported_namespaces(self) -> Iterator[Tuple[Node, str]]:
+        for node, _ in self.get_code_sections():
+            if self.is_aspx_import_directive(node):
+                namespace = self.get_aspx_namespace(node)
+                if namespace is not None:
+                    yield node, namespace
+
+    def get_template_namespaces(self) -> Iterator[Tuple[Optional[Node], str]]:
+        for namespace in self.language_toolkit.get_default_namespaces(True):
+            yield None, namespace
+        for node, namespace in self.get_imported_namespaces():
+            yield node, namespace
+
+    def is_c_sharp(self, node: Node) -> bool:
+        return bool(
+            re.match(
+                r'@ .*Page Language\s*=\s*"C#".*'.encode(),
+                self.get_byte_range(node),
+                re.IGNORECASE,
+            )
+        )
+
+    def is_aspx_import_directive(self, node: Node) -> bool:
+        return bool(
+            re.match(
+                r"@\s*Import Namespace=".encode(),
+                self.get_byte_range(node),
+                re.IGNORECASE,
+            )
+        )
+
+    def get_aspx_namespace(self, node: Node) -> Optional[str]:
+        match = re.search(
+            r'@\s*Import namespace="(.*?)"'.encode(),
+            self.get_byte_range(node),
+            re.IGNORECASE,
+        )
+        return match.group(1).decode() if match is not None else None
+
+
+class TreeSitterHTMLEngine(TreeSitterBaseEngine):
+    query: HTMLQueryBinding
+    namespaces: set[str]
+
+    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(node)
+
+    def get_identified_scripts(self) -> Iterator[Tuple[Node, str]]:
+        for node, _ in self.get_scripts():
+            for content_node, _ in self.get_script_contents(node):
+                yield content_node, self.identify_language(node)
+
+    def get_script_contents(self, node: Node) -> Iterator[Tuple[Node, str]]:
+        return self.query.script_content.captures(node)
+
+    def get_parsed_code_sections(self) -> Iterator[TreeSitterExtractorEngine]:
+        for node, language in self.get_identified_scripts():
+            # TODO: support JS
+            if language == LANG_CS:
+                yield TreeSitterExtractorEngine(language, self.get_byte_range(node), node.start_byte, self.namespaces)
+
+    def identify_language(self, node: Node) -> str:
+        for attribute_node, _ in self.get_attributes(node):
+            if self.is_server_side_c_sharp(attribute_node):
+                return LANG_CS
+        return LANG_JS
+
+    def is_server_side_c_sharp(self, node: Node) -> bool:
+        return len(re.findall(r'runat\s*=\s*"server"'.encode(), self.get_byte_range(node))) > 0