fix: use new TS API (#2)

Update and pin to current TreeSitter Python implementation.

Author: 0xricksanchez

__init__.py

@@ -10,23 +10,16 @@
 from typing import Dict, List, Optional, Tuple
 import json
 import re
-import struct
 import subprocess
 import html
 
 from tree_sitter import Language, Parser, Query
+import tree_sitter_c
 import emoji
 
 from binaryninja import *
 
 
-# directory path to the current script
-CURRENT_DIR = Path(__file__).parent
-
-# tree-sitter files
-TREE_SITTER_C = CURRENT_DIR / "tree-sitter-c"
-TREE_SITTER_LIB = CURRENT_DIR / "build" / "tree-sitter-c.so"
-
 # query to search tree-sitter's syntax tree for illegal identifier annotations
 FUNC_ANNOT_QUERY_STR = """
 (function_definition
@@ -249,15 +242,9 @@ def __init__(self, bv):
             self.bv, disas_settings
         )
 
-        # initialize tree-sitter
-        if not TREE_SITTER_LIB.is_file():
-            Language.build_library(str(TREE_SITTER_LIB), [str(TREE_SITTER_C)])
-            if not TREE_SITTER_LIB.is_file():
-                raise Exception("Failed to build tree-sitter lib")
-
-        c_language = Language(str(TREE_SITTER_LIB), "c")
+        c_language = Language(tree_sitter_c.language())
         self.parser = Parser()
-        self.parser.set_language(c_language)
+        self.parser.language = c_language
         self.func_annot_query = c_language.query(FUNC_ANNOT_QUERY_STR)
 
     def run(self):
@@ -284,11 +271,10 @@ def run(self):
             ]
 
             # generate temp files for pseudo C code and semgrep findings
-            with NamedTemporaryFile(
-                suffix=".c", mode="w+"
-            ) as pseudo_c_out, NamedTemporaryFile(
-                suffix=".json", mode="w+"
-            ) as semgrep_results_out:
+            with (
+                NamedTemporaryFile(suffix=".c", mode="w+") as pseudo_c_out,
+                NamedTemporaryFile(suffix=".json", mode="w+") as semgrep_results_out,
+            ):
                 # serialise pseudo C contents
                 for gv in global_vars:
                     pseudo_c_out.write(f"{gv}\n")
@@ -446,10 +432,12 @@ def remove_function_annotations(self, src: str) -> str:
         captures = self.func_annot_query.captures(tree.root_node)
         src_list = list(src)
 
-        for node, _ in captures:
-            # replace each annotation with the empty string
-            for i in range(node.start_byte, node.end_byte):
-                src_list[i] = ""
+        # Process all nodes from all capture groups
+        for capture_list in captures.values():
+            for node in capture_list:
+                # replace each annotation with the empty string
+                for i in range(node.start_byte, node.end_byte):
+                    src_list[i] = ""
 
         # reconstruct source code from the list
         return "".join(src_list)

plugin.json

@@ -1,12 +1,16 @@
 {
     "pluginmetadataversion": 2,
     "name": "semgrep-bn",
-    "version": "1.0",
+    "version": "1.0.1",
     "author": "Samman Palihapitiya",
-    "type": [ "helper" ],
+    "type": [
+        "helper"
+    ],
     "description": "Search code with Semgrep",
     "longdescription": "",
-    "api": [ "python3" ],
+    "api": [
+        "python3"
+    ],
     "platforms": [
         "Darwin",
         "Linux",

requirements.txt

@@ -1,3 +1,4 @@
 emoji
-tree-sitter
-semgrep
+tree-sitter==0.23.2
+tree-sitter-c==0.23.4
+semgrep