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) == (other.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,10 +405,17 @@ 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"
-VALID_FORMAT = (FORMAT_PE, FORMAT_ELF, FORMAT_DOTNET)
+FORMAT_SCRIPT = "script"
+VALID_FORMAT = (FORMAT_PE, FORMAT_ELF, FORMAT_DOTNET, FORMAT_SCRIPT)
 # internal only, not to be used in rules
 FORMAT_AUTO = "auto"
 FORMAT_SC32 = "sc32"

capa/features/extractors/common.py

@@ -9,9 +9,21 @@
 import capa.features
 import capa.features.extractors.elf
 import capa.features.extractors.pefile
-from capa.features.common import OS, FORMAT_PE, FORMAT_ELF, OS_WINDOWS, FORMAT_FREEZE, Arch, Format, String, Feature
+from capa.features.common import (
+    OS,
+    FORMAT_PE,
+    FORMAT_ELF,
+    OS_WINDOWS,
+    FORMAT_FREEZE,
+    FORMAT_SCRIPT,
+    Arch,
+    Format,
+    String,
+    Feature,
+)
 from capa.features.freeze import is_freeze
 from capa.features.address import NO_ADDRESS, Address, FileOffsetAddress
+from capa.features.extractors.ts.autodetect import is_script
 
 logger = logging.getLogger(__name__)
 
@@ -34,6 +46,8 @@ def extract_format(buf) -> Iterator[Tuple[Feature, Address]]:
         yield Format(FORMAT_ELF), NO_ADDRESS
     elif is_freeze(buf):
         yield Format(FORMAT_FREEZE), NO_ADDRESS
+    elif is_script(buf):
+        yield Format(FORMAT_SCRIPT), NO_ADDRESS
     else:
         # we likely end up here:
         #  1. handling a file format (e.g. macho)

capa/features/extractors/script.py

@@ -0,0 +1,41 @@
+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
+
+# Can be used to instantiate tree_sitter Language objects (see ts/query.py)
+LANG_CS = "c_sharp"
+LANG_HTML = "html"
+LANG_JS = "javascript"
+LANG_PY = "python"
+LANG_TEM = "embedded_template"
+
+EXT_ASPX = ("aspx", "aspx_")
+EXT_CS = ("cs", "cs_")
+EXT_HTML = ("html", "html_")
+EXT_PY = ("py", "py_")
+
+
+LANGUAGE_FEATURE_FORMAT = {
+    LANG_CS: "C#",
+    LANG_HTML: "HTML",
+    LANG_JS: "JavaScript",
+    LANG_PY: "Python",
+    LANG_TEM: "Embedded Template",
+}
+
+
+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_FEATURE_FORMAT[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

capa/features/extractors/ts/autodetect.py

@@ -0,0 +1,65 @@
+from typing import Optional
+
+from tree_sitter import Node, Tree, Parser, Language
+
+from capa.features.extractors.script import EXT_CS, EXT_PY, LANG_CS, LANG_PY, EXT_ASPX, EXT_HTML, LANG_TEM, LANG_HTML
+from capa.features.extractors.ts.query import TS_LANGUAGES
+
+
+def is_script(buf: bytes) -> bool:
+    try:
+        return bool(get_language_ts(buf))
+    except ValueError:
+        return False
+
+
+def _parse(ts_language: Language, buf: bytes) -> Optional[Tree]:
+    try:
+        parser = Parser()
+        parser.set_language(ts_language)
+        return parser.parse(buf)
+    except ValueError:
+        return None
+
+
+def _contains_errors(ts_language, node: Node) -> bool:
+    return ts_language.query("(ERROR) @error").captures(node)
+
+
+def get_language_ts(buf: bytes) -> str:
+    for language, ts_language in TS_LANGUAGES.items():
+        tree = _parse(ts_language, buf)
+        if tree and not _contains_errors(ts_language, tree.root_node):
+            return language
+    raise ValueError("failed to parse the language")
+
+
+def get_template_language_ts(buf: bytes) -> str:
+    for language, ts_language in TS_LANGUAGES.items():
+        if language in [LANG_TEM, LANG_HTML]:
+            continue
+        tree = _parse(ts_language, buf)
+        if tree and not _contains_errors(ts_language, tree.root_node):
+            return language
+    raise ValueError("failed to parse the language")
+
+
+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
+    if path.endswith(EXT_PY):
+        return LANG_PY
+    raise ValueError(f"{path} has an unrecognized or an unsupported extension.")
+
+
+def get_language(path: str) -> str:
+    try:
+        with open(path, "rb") as f:
+            buf = f.read()
+        return get_language_ts(buf)
+    except ValueError:
+        return get_language_from_ext(path)

capa/features/extractors/ts/build.py

@@ -0,0 +1,15 @@
+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",
+    "vendor/tree-sitter-python",
+]
+
+
+class TSBuilder:
+    def __init__(self):
+        Language.build_library(build_dir, languages)