import os
import re


class SmaliParser:
    """
    smali文件解析器
    """

    def __init__(self, _root):
        """
        构造函数
        :param _root: 传入需要解析smali文件的根目录
        """
        self.root = _root
        self.methods = []

    def __get_files_path(self, ) -> [str]:
        """
        输入文件根目录，返回该目录下所有smali文件的路径(包含递归子目录)
        :return: [str,...], 所有smali文件目录
        """
        fs = []
        for root, dirs, files in os.walk(self.root):
            for f in files:
                if f.endswith('.smali'):
                    fs.append(os.path.join(root, f))
        return fs

    def __parse_smali(self, file: str):
        """
        解析samli文件,返回method,以及其中所调用的函数
        :param file: str, smali文件路径
        :return: [
                {'name': name, 'args': args, 'return': ret,
                'calls': [{'invokeType': invoke, 'package': package, 'name': name, 'args': args, 'return': ret},...]
                }
                ,...]
        """
        def format_name():
            pass
        # def
        f = open(file, 'r')
        lines = f.readlines()
        f.close()
        methods = []  # 存储所有的method
        calls = []  # 每个method的调用函数
        flag = False  # 标识是否在.method
        for line in lines:
            try:
                if re.search('\.method', line) is not None:
                    flag = True
                    # .method public postToMainThread(Ljava/lang/Runnable;)V
                    #                   name                args            ret
                    name, args, ret = re.split(r'[\(|\)]', line.split(' ')[-1])  # 解析method属性
                    methods.append({'name': name, 'args': args, 'return': ret, 'calls': []})
                elif flag and re.search('invoke-', line) is not None:
                    while line[0] == ' ' or line[0] == '\t' or line[0] == '\r' or line[0] == '\f':  # 移除开始空格符号
                        line = line[1:]
                    # invoke-virtual {v0, p1}, Landroid/arch/TaskExecutor;->postToMainThread(Ljava/lang/Runnable;)V
                    #       invoke                     package                     name            args            ret
                    invoke, _, attribute = re.split(r'[{|}]', line)
                    invoke = invoke[:-1]  # 去掉最后空格
                    attribute = attribute[2:]  # 去掉前面,和空格
                    package, attribute = attribute.split(';->')
                    name, args, ret = re.split(r'[\(|\)]', attribute)  # 解析调用函数属性
                    calls.append({'invokeType': invoke, 'package': package, 'name': name, 'args': args, 'return': ret})
                elif flag and re.search('.end method', line) is not None:
                    methods[-1]['calls'] = calls  # 当前method完毕，将calls更新到最后一个method，清空calls
                    calls = []
            except:
                print("\033[31mParse smali file error! In [" + file + "],line [" + line + "] parse error" + "\033[0m")
        return methods

    def get_results(self):
        files = self.__get_files_path()
        for f in files:
            t = self.__parse_smali(f)
            for j in t:
                self.methods.append(j)
        return self.methods


if __name__ == '__main__':
    ms = SmaliParser('E:\\apktooltest\\3\\smali').get_results()
