Simple test code to scan malicious py scripts

tree_sitter_py_scan.py

"""
Code parser with tree sitter
`sudo pip3 install tree_sitter`
clone parser for each programming language (same dir with code py) `git clone https://github.com/tree-sitter/tree-sitter-python`
create test code like eval(base64.decode(<base64_text>))
"""


from tree_sitter import Language, Parser

# Compile sources to parser. Compile time only
Language.build_library('build/parser.so', ['tree-sitter-c', 'tree-sitter-python'])
# Init Parser
C_LANG = Language('build/parser.so', 'c')
PY_LANG = Language('build/parser.so', 'python')


def handle_file(path):
	with open(path) as f:
		return f.read()


def analysis_function(name, argv):
	if name.text.decode() in ("eval", "exec"):
		print(f"[!] Dangerous function {name.text.decode()} Code execution")
	else:
		# Calculate buffer?
		pass


def parse_code(cursor):
	if cursor.type == "call":
		analysis_function(*cursor.children)


def parse_script(script_data):
	if script_data:
		parser = Parser()
		parser.set_language(PY_LANG)
		tree = parser.parse(bytes(script_data, "utf-8"))
		cursor = tree.walk()
		cursor.goto_first_child()
		while True:
			if cursor.node.type == "function_definition":
				parse_code(cursor.node.children[2])
			elif cursor.node.type == "expression_statement":
				# In Py code, there's no need to create main function
				parse_code(cursor.node.children[0])
			if not cursor.goto_next_sibling():
				return


data = handle_file("/tmp/encoded_revshell.py")
parse_script(data)