[+]support basic_block

[+]support Block Object
[+]change the method of find method

cfg_generator.py

@@ -1,42 +1,126 @@
 from models.inner_instruction import *
 from models.cfg import *
 
+def generate_cfg_block(block, info_provider, class_name, method_name, recursive=False):
 
-def generate_cfg(method, info_provider, recursive=False):
-    print('Current generate CFG of method: %s in class: %s' % (method.method_name, method.class_name))
+    # cfg_blocks = []
+    cfg_block = CFGBlock(block.identify)
 
-    cfg_name = method.class_name + ': ' + method.method_name
-    cfg = CFG(cfg_name)
-    for i in range(len(method.instructions)):
-        instruction = method.instructions[i]
-        if instruction.goto_insns:
+    # wait_for_follow = []
+
+    for i in range((len(block.instructions))):
+        instruction = block.instructions[i]
+        # if cfg_block is None:
+        #     cfg_block = CFGBlock(hex(instruction.address))
+        #     if len(wait_for_follow) > 0:
+        #         for fol_block in wait_for_follow:
+        #             fol_block.goto_block(cfg_block)
+        #         wait_for_follow = []
+
+        if instruction.goto_insns:  # 如果有调用函数
             basic_info, imp_name = instruction.goto_insns
             if basic_info == '$Function':  # function
                 if filter_oc_function(imp_name):
                     continue
-                if recursive:
+                if recursive and basic_info != class_name and method_name != method_name:  # 如果需要更进一步解析（防止递归）
                     recursive_function = info_provider(basic_info, imp_name)
                     if recursive_function is not None:
                         recursive_cfg = generate_cfg(recursive_function, info_provider, True)
-                        for recursive_cfg_node in recursive_cfg.nodes:
-                            cfg.add_node(recursive_cfg_node)
+                        cfg_block.add_node(recursive_cfg)
+                        # cfg_block.goto_block(recursive_cfg.entry)  # 该块进入调用的函数
+                        # cfg_blocks.append(cfg_block)
+
+                        # for rec_block in recursive_cfg.all_blocks:  # 将该函数的块都加到当前 CFG 中
+                        #     cfg_blocks.append(rec_block)
+                        #     if rec_block.out:
+                        #         wait_for_follow.append(rec_block)
+
+                        # cfg_block = None
+                        # cfg_block.add_node(recursive_cfg)
+                        # cfg_blocks.append(cfg_block)
                         continue
                 cfg_node = CFGNode(CFGNodeTypeFunction)
                 cfg_node.function_name = imp_name
-            else:  # method
+            else:
                 if recursive:
                     recursive_method = info_provider(basic_info, imp_name)
                     if recursive_method is not None:
                         recursive_cfg = generate_cfg(recursive_method, info_provider, True)
-                        for recursive_cfg_node in recursive_cfg.nodes:
-                            cfg.add_node(recursive_cfg_node)
+                        cfg_block.add_node(recursive_cfg)
                         continue
                 cfg_node = CFGNode(CFGNodeTypeMethod)
                 cfg_node.class_name = basic_info
                 cfg_node.method_name = imp_name
-            cfg.add_node(cfg_node)
+            cfg_block.add_node(cfg_node)
+
+    # if cfg_block is not None:
+    #     cfg_blocks.append(cfg_block)
+
+    return cfg_block
+
+
+def generate_cfg(method, info_provider, recursive=False):
+    print('Current generate CFG of method: %s in class: %s' % (method.method_name, method.class_name))
+
+    wait_blocks_queue = []
+
+    cfg_name = method.class_name + ': ' + method.method_name
+    cfg = CFG(cfg_name)
+
+    wait_blocks_queue.append(method.entry_block)
+
+    while len(wait_blocks_queue) > 0:
+        block = wait_blocks_queue[0]
+        wait_blocks_queue = wait_blocks_queue[1:]
+
+        cfg_block = generate_cfg_block(block, info_provider, method.class_name,
+                                       method.method_name, recursive)
+        cfg.add_block(cfg_block)
+        if cfg.entry is None:
+            cfg.entry = cfg_block
+
+        if not block.is_return:  # 含有 ret 的基本块应该肯定不会有 followed 的块
+            if block.jump_to_block is not None and block.jump_to_block in method.all_blocks:
+                cfg_block.goto_block(block.jump_to_block)
+                if cfg.get_block(block.jump_to_block) is None:
+                    wait_blocks_queue.append(method.all_blocks[block.jump_to_block])
+
+            if block.jump_condition and block.next_block is not None:
+                cfg_block.goto_block(block.next_block)
+                if cfg.get_block(block.next_block) is None:
+                    wait_blocks_queue.append(method.all_blocks[block.next_block])
     return cfg
 
+    # for i in range(len(method.instructions)):
+    #     instruction = method.instructions[i]
+    #     if instruction.goto_insns:
+    #         basic_info, imp_name = instruction.goto_insns
+    #         if basic_info == '$Function':  # function
+    #             if filter_oc_function(imp_name):
+    #                 continue
+    #             if recursive:
+    #                 recursive_function = info_provider(basic_info, imp_name)
+    #                 if recursive_function is not None:
+    #                     recursive_cfg = generate_cfg(recursive_function, info_provider, True)
+    #                     for recursive_cfg_node in recursive_cfg.nodes:
+    #                         cfg.add_node(recursive_cfg_node)
+    #                     continue
+    #             cfg_node = CFGNode(CFGNodeTypeFunction)
+    #             cfg_node.function_name = imp_name
+    #         else:  # method
+    #             if recursive:
+    #                 recursive_method = info_provider(basic_info, imp_name)
+    #                 if recursive_method is not None:
+    #                     recursive_cfg = generate_cfg(recursive_method, info_provider, True)
+    #                     for recursive_cfg_node in recursive_cfg.nodes:
+    #                         cfg.add_node(recursive_cfg_node)
+    #                     continue
+    #             cfg_node = CFGNode(CFGNodeTypeMethod)
+    #             cfg_node.class_name = basic_info
+    #             cfg_node.method_name = imp_name
+    #         cfg.add_node(cfg_node)
+    # return cfg
+
 
 def filter_oc_function(function):
     if function.startswith('_objc_'):

interpreters/inner_Interpreter.py

@@ -4,7 +4,8 @@
 from capstone.arm64 import *
 
 SELF_POINTER = -0x1000000
-SUPER_POINTER = -0x2000000
+CURRENT_SELECTOR = -0x2000000
+
 
 class Register:
 
@@ -50,19 +51,30 @@ def clear(self):
 
 class Interpreter:
 
-    def __init__(self, memory_provider=None, handle_strange_add=None):
+    def __init__(self, memory_provider=None, handle_strange_add=None, parameters=[]):
         self.gen_regs = [Register(i) for i in range(31)]
         self.float_regs = [FloatRegister(i) for i in range(32)]
-        self.gen_regs[0].value = SELF_POINTER
+
         self.wzr = Register(-1)
         self.xzr = Register(-1)
         self.wsp = Register(-1)
         self.sp = Register(-1)
         self.pc = Register(-1)
-        self.memory = {hex(0-0x30): SUPER_POINTER}
+        self.memory = {}
+        # self.memory = {hex(0-0x30): SUPER_POINTER}  # 父指针好像在
         self.memory_provider = memory_provider
         self.handle_strange_add = handle_strange_add
 
+        if len(parameters) <= 4:
+            for i in range(len(parameters)):
+                self.gen_regs[i].value = parameters[i]
+        else:
+            for i in range(4):
+                self.gen_regs[i].value = parameters[i]
+            # 超过 4 个参数再说
+            # for i in range(len(parameters) - 4):
+                # self.memory[hex()]
+
         # Jump related
         self.compare_flag = 0  # 0 is equal and -1 is small and 1 is bigger
         self.should_jump = False
@@ -94,7 +106,7 @@ def clear_regs(self):
         self.wsp.clear()
         self.sp.clear()
         self.pc.clear()
-        self.memory = {hex(0-0x30):SUPER_POINTER}
+        # self.memory = {hex(0-0x30):SUPER_POINTER}
 
     def interpret_code(self, codes, begin=0, end=-1):
         i = begin
@@ -264,7 +276,7 @@ def handle_add(self, insn):
         if insn.operands[1].type == ARM64_OP_REG and insn.operands[2].type == ARM64_OP_REG:
             reg_name = insn.reg_name(insn.operands[1].reg)
             register = self.get_register(reg_name)
-            if register.value < 0:
+            if register.value < 0:  # 在取 ivar 的时候，会遇到这种问题，因为现在对于 SELF 指针的定义为一个负数的常量
                 reg_name_2 = insn.reg_name(insn.operands[2].reg)
                 register_2 = self.get_register(reg_name_2)
                 dest = insn.operands[0]

models/cfg.py

@@ -1,41 +1,125 @@
+from graphviz import Digraph, Graph
+
 CFGNodeTypeFunction = 0
 CFGNodeTypeMethod = 1
 
 
 class CFG:
 
-    def __init__(self, name=''):
+    def __init__(self, name='', entry=None):
+        self.name = name
+        self.entry = entry
+        self.outs = []
+        self.all_blocks = []
+
+    def add_block(self, block):
+        self.all_blocks.append(block)
+        if block.out:
+            self.outs.append(block)
+
+    def get_block(self, name):
+        for block in self.all_blocks:
+            if block.name == name:
+                return block
+        return None
+    # if the data flow between the no.0 node and no.1 node
+    # the index should be 0
+    # def modify_data_flow(self, index, data_flow):
+    #     node_count = len(self.nodes)
+    #     if node_count - 1 > index:
+    #         if index < len(self.data_flows):
+    #             self.data_flows[index] = data_flow
+    #         else:
+    #             for i in range(index):
+    #                 self.data_flows.append(None)
+    #             self.data_flows.append(data_flow)
+    #     else:
+            # print('Data flow index error!')
+
+    def describe(self):
+        for block in self.all_blocks:
+            block.describe()
+
+    def graphviz_obj(self):
+
+        cfg_view = Digraph(self.name)
+        for block in self.all_blocks:
+            for node in block.nodes:
+                if type(node) == CFGNode:
+                    node_name = block.name + str(block.nodes.index(node))
+                    node_label = node.describe(False)
+                    cfg_view.node(node_name, node_label)
+                elif type(node) == CFG:
+                    pass
+
+        return cfg_view
+
+        # block_nodes = {}
+        # cfg_view = Digraph(self.name)
+        # # cfg_view.node('entry', 'entry')
+        #
+        # for block in self.all_blocks:
+        #     block_view = Digraph('cluster' + block.name)
+        #     if len(block.nodes) == 0:
+        #         node_name = 'node' + block.name
+        #         block_view.node(node_name, '', shape='plaintext')
+        #         block_nodes[block.name] = (node_name, node_name)
+        #     else:
+        #         start_name = None
+        #         end_name = None
+        #         before_name = None
+        #         for node in block.nodes:
+        #             if type(node) == CFG:
+        #                 pass
+        #                 # recursive_cfg = node.graphviz_obj()
+        #                 # block_view.subgraph(recursive_cfg)
+        #
+        #             elif type(node) == CFGNode:
+        #                 node_name = block.name + str(block.nodes.index(node))
+        #                 node_label = node.describe(False)
+        #                 block_view.node(node_name, node_label, shape='box')
+        #                 if before_name is not None:
+        #                     block_view.edge(before_name, node_name)
+        #                 before_name = node_name
+        #                 if start_name is None:
+        #                     start_name = node_name
+        #                 if block == self.all_blocks[-1]:
+        #                     end_name = end_name
+        #         block_nodes[block.name] = (start_name, end_name)
+        #
+        #     cfg_view.subgraph(block_view)
+        #
+        # # cfg_view.edge('entry', 'cluster' + self.entry.name)
+        #
+        # for block in self.all_blocks:
+        #     _, first_name = block_nodes[block.name]
+        #     for follow in block.follow_blocks:
+        #         second_name, _ = block_nodes[follow]
+        #         cfg_view.edge(first_name, second_name)
+        # return cfg_view
+
+    def view(self):
+        self.graphviz_obj().view()
+
+class CFGBlock:
+
+    def __init__(self, name):
         self.name = name
-        self.nodes = []
-        self.data_flows = []
+        self.out = False  # if this block contains `ret` instruction
+        self.nodes = []  # node 包括 node 或者 cfg
+        self.follow_blocks = []  # the blocks follow this block (name)
 
     def add_node(self, node):
         self.nodes.append(node)
 
-    # if the data flow between the no.0 node and no.1 node
-    # the index should be 0
-    def modify_data_flow(self, index, data_flow):
-        node_count = len(self.nodes)
-        if node_count - 1 > index:
-            if index < len(self.data_flows):
-                self.data_flows[index] = data_flow
-            else:
-                for i in range(index):
-                    self.data_flows.append(None)
-                self.data_flows.append(data_flow)
-        else:
-            print('Data flow index error!')
+    def goto_block(self, block):
+        self.follow_blocks.append(block)
 
     def describe(self):
-        for i in range(1, len(self.nodes) - 1):
+        for i in range(len(self.nodes) - 1):
             self.nodes[i].describe()
-            if (i - 1 < len(self.data_flows) and
-                self.data_flows[i - 1] is not None and
-                not self.data_flows[i - 1].isEmpty):
-                self.data_flows[i - 1].describe()
-            else:
-                print('||')
-                print('\/')
+            print('||')
+            print('\/')
         if len(self.nodes) > 0:
             self.nodes[-1].describe()
 
@@ -54,11 +138,14 @@ def node_info(self):
         else:
             return self.class_name, self.method_name
 
-    def describe(self):
+    def describe(self, verbose=True):
         if self.type == CFGNodeTypeFunction:
             print('<%s>' % self.function_name)
+            describe = '<' + self.function_name + '>'
         else:
             print('<%s: %s>' % (self.class_name, self.method_name))
+            describe = '<' + self.class_name + ': ' + self.method_name + '>'
+        return describe
 
 
 class CFGDataFlow:

models/class_storage.py

@@ -46,6 +46,18 @@ def __init__(self, name):
         self.arguments_type = []  # empty means no argument
 
 
+BlockMethodTypeStack = 0
+BlockMethodTypeGlobal = 1
+BlockMethodTypeMalloc = 2
+
+
+class BlockMethodData:
+
+    def __init__(self, type=BlockMethodTypeStack):
+        self.type = type
+        self.invoke = 0
+
+
 MethodDataTypeClass = 0
 MethodDataTypeInstance = 1