Example/ctype: clean code

1 files changed, 110 insertions, 129 deletions

diff --git a/example/ida/ctype_propagation.py b/example/ida/ctype_propagation.py
index e9986522..95fef75f 100644
--- a/example/ida/ctype_propagation.py
+++ b/example/ida/ctype_propagation.py
@@ -28,12 +28,11 @@ class TypePropagationForm(Form):
         self.ira = ira
         self.stk_unalias_force = False
 
-        default_types_info = r"""ExprId("RDX", 64): char *
-"""
+        default_types_info = r"""ExprId("RDX", 64): char *"""
         archs = ["AMD64_unk", "X86_32_unk"]
 
         Form.__init__(self,
-r"""BUTTON YES* Launch
+                      r"""BUTTON YES* Launch
 BUTTON CANCEL NONE
 Dependency Graph Settings
 <##Header file :{headerFile}>
@@ -41,15 +40,15 @@ Dependency Graph Settings
 <Types informations:{strTypesInfo}>
 <Unalias stack:{rUnaliasStack}>{cMethod}>
 """, {
-    'headerFile': Form.FileInput(swidth=20, open=True),
-    'cbReg': Form.DropdownListControl(
-        items=archs,
-        readonly=False,
-        selval=archs[0]),
-    'strTypesInfo': Form.MultiLineTextControl(text=default_types_info,
-                                              flags=Form.MultiLineTextControl.TXTF_FIXEDFONT),
-    'cMethod': Form.ChkGroupControl(("rUnaliasStack",)),
-    })
+                          'headerFile': Form.FileInput(swidth=20, open=True),
+                          'cbReg': Form.DropdownListControl(
+                              items=archs,
+                              readonly=False,
+                              selval=archs[0]),
+                          'strTypesInfo': Form.MultiLineTextControl(text=default_types_info,
+                                                                    flags=Form.MultiLineTextControl.TXTF_FIXEDFONT),
+                          'cMethod': Form.ChkGroupControl(("rUnaliasStack",)),
+                      })
         self.Compile()
 
     @property
@@ -82,38 +81,6 @@ def get_types_mngr(headerFile):
     return types_mngr
 
 
-
-# add stack alignment simpl
-def simp_stack_align(expr_simp, expr):
-    if not expr.is_op("&"):
-        return expr
-    if len(expr.args) != 2:
-        return expr
-    if not expr.args[1].is_int(0xFFFFFFF0):
-        return expr
-    if expr.args[0] != ir_arch.arch.regs.ESP_init:
-        return expr
-    return expr.args[0]
-
-expr_simp.enable_passes({m2_expr.ExprOp: [simp_stack_align]})
-
-# Init
-machine = guess_machine()
-mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
-
-bs = bin_stream_ida()
-mdis = dis_engine(bs, dont_dis_nulstart_bloc=True)
-ir_arch = ira(mdis.symbol_pool)
-
-# Get the current function
-func = idaapi.get_func(ScreenEA())
-addr = func.startEA
-blocks = mdis.dis_multibloc(addr)
-# Generate IR
-for block in blocks:
-    ir_arch.add_bloc(block)
-
-
 class MyCTypeAnalyzer(CTypeAnalyzer):
     allow_none_result = True
 
@@ -127,21 +94,93 @@ class MyCHandler(CHandler):
     exprToAccessC_cls = MyExprToAccessC
 
 
-# Get settings
-settings = TypePropagationForm(ir_arch)
-ret = settings.Execute()
-if ret:
-
-    ir_arch = ir_arch
+class TypePropagationEngine(SymbExecCType):
+
+    def __init__(self, ir_arch, state):
+        mychandler = MyCHandler(types_mngr, state.infos_types)
+        super(TypePropagationEngine, self).__init__(ir_arch,
+                                                    state.symbols,
+                                                    state.infos_types,
+                                                    mychandler)
+
+
+class SymbExecCTypeFix(SymbExecCType):
+
+    def emulbloc(self, irb, step=False):
+        """
+        Symbolic execution of the @irb on the current state
+        @irb: irblock instance
+        @step: display intermediate steps
+        """
+        offset2cmt = {}
+        for assignblk in irb.irs:
+            instr = assignblk.instr
+            tmp_rw = assignblk.get_rw()
+            for dst, src in assignblk.iteritems():
+                for arg in set(instr.args).union(set([src])):
+                    if arg in tmp_rw and arg not in tmp_rw.values():
+                        continue
+                    objc = self.eval_expr(arg)
+                    if objc is None:
+                        continue
+                    if self.is_type_offset(objc):
+                        continue
+                    offset2cmt.setdefault(instr.offset, set()).add(
+                        "%s: %s" % (arg, str(objc)))
+            self.eval_ir(assignblk)
+
+        for offset, value in offset2cmt.iteritems():
+            idc.MakeComm(offset, '\n'.join(value))
+
+        return self.eval_expr(self.ir_arch.IRDst)
+
+
+class CTypeEngineFixer(SymbExecCTypeFix):
+
+    def __init__(self, ir_arch, state):
+        mychandler = MyCHandler(types_mngr, state.infos_types)
+        super(CTypeEngineFixer, self).__init__(ir_arch,
+                                               state.symbols,
+                                               state.infos_types,
+                                               mychandler)
+
+
+def add_state(ir_arch, todo, states, addr, state):
+    addr = ir_arch.get_label(addr)
+    if addr not in states:
+        states[addr] = state
+        todo.add(addr)
+    else:
+        todo.add(addr)
+        states[addr] = states[addr].merge(state)
+
+
+def analyse_function():
+
+    # Init
+    machine = guess_machine()
+    mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
+
+    bs = bin_stream_ida()
+    mdis = dis_engine(bs, dont_dis_nulstart_bloc=True)
+    ir_arch = ira(mdis.symbol_pool)
+
+    # Get the current function
+    func = idaapi.get_func(ScreenEA())
+    addr = func.startEA
+    blocks = mdis.dis_multibloc(addr)
+    # Generate IR
+    for block in blocks:
+        ir_arch.add_bloc(block)
 
+    # Get settings
+    settings = TypePropagationForm(ir_arch)
+    ret = settings.Execute()
+    if not ret:
+        return
 
     types_mngr = get_types_mngr(settings.headerFile.value)
     mychandler = MyCHandler(types_mngr, {})
-
-    #print 'Expr', settings.iStr1.value
-    #print 'Type', settings.iStr2.value
-    #print 'Header', settings.headerFile.value
-    print 'InfoTypes', settings.strTypesInfo.value
     infos_types = {}
     for line in settings.strTypesInfo.value.split('\n'):
         if not line:
@@ -149,10 +188,12 @@ if ret:
         expr_str, ctype_str = line.split(':')
         expr_str, ctype_str = expr_str.strip(), ctype_str.strip()
         expr = str_to_expr(expr_str)
-        ast = mychandler.type_analyzer.types_mngr.types_ast.parse_c_type(ctype_str)
-        ctype = mychandler.type_analyzer.types_mngr.types_ast.ast_parse_declaration(ast.ext[0])
+        ast = mychandler.type_analyzer.types_mngr.types_ast.parse_c_type(
+            ctype_str)
+        ctype = mychandler.type_analyzer.types_mngr.types_ast.ast_parse_declaration(ast.ext[
+                                                                                    0])
         objc = types_mngr.get_objc(ctype)
-        print '='*20
+        print '=' * 20
         print expr, objc
         infos_types[expr] = objc
 
@@ -163,33 +204,15 @@ if ret:
     first_block = blocks.label2block(lbl_real_start)
 
     assignblk_head = AssignBlock([ExprAff(ir_arch.IRDst, ExprId(lbl_real_start, ir_arch.IRDst.size)),
-                                  ExprAff(ir_arch.sp, ir_arch.arch.regs.regs_init[ir_arch.sp])
-                                 ], first_block.lines[0])
+                                  ExprAff(
+                                      ir_arch.sp, ir_arch.arch.regs.regs_init[ir_arch.sp])
+                                  ], first_block.lines[0])
     irb_head = IRBlock(lbl_head, [assignblk_head])
     ir_arch.blocks[lbl_head] = irb_head
     ir_arch.graph.add_uniq_edge(lbl_head, lbl_real_start)
 
-    class Engine(SymbExecCType):
-        def __init__(self, state):
-            mychandler = MyCHandler(types_mngr, state.infos_types)
-            super(Engine, self).__init__(ir_arch,
-                                         state.symbols,
-                                         state.infos_types,
-                                         mychandler)
-
-
-    def add_state(todo, states, addr, state):
-        addr = ir_arch.get_label(addr)
-        if addr not in states:
-            states[addr] = state
-            todo.add(addr)
-        else:
-            todo.add(addr)
-            states[addr] = states[addr].merge(state)
-
-
-    state = Engine.StateEngine(infos_types, infos_types)
-    states = {lbl_head:state}
+    state = TypePropagationEngine.StateEngine(infos_types, infos_types)
+    states = {lbl_head: state}
     todo = set([lbl_head])
     done = set()
 
@@ -199,7 +222,7 @@ if ret:
         if (lbl, state) in done:
             continue
         done.add((lbl, state))
-        symbexec_engine = Engine(state)
+        symbexec_engine = TypePropagationEngine(ir_arch, state)
 
         get_block(ir_arch, mdis, lbl)
 
@@ -211,56 +234,14 @@ if ret:
         for son in sons:
             if son.offset is None:
                 continue
-            add_state(todo, states, son.offset, symbexec_engine.get_state())
-
-
-    class SymbExecCTypeFix(SymbExecCType):
-        def emulbloc(self, irb, step=False):
-            """
-            Symbolic execution of the @irb on the current state
-            @irb: irblock instance
-            @step: display intermediate steps
-            """
-            offset2cmt = {}
-            for assignblk in irb.irs:
-                instr = assignblk.instr
-                tmp_rw = assignblk.get_rw()
-                for dst, src in assignblk.iteritems():
-                    for arg in set(instr.args).union(set([src])):
-                        if arg in tmp_rw and arg not in tmp_rw.values():
-                            continue
-                        objc = self.eval_expr(arg)
-                        if objc is None:
-                            continue
-                        if self.is_type_offset(objc):
-                            continue
-                        offset2cmt.setdefault(instr.offset, set()).add("%s: %s" % (arg, str(objc)))
-                self.eval_ir(assignblk)
-
-            for offset, value in offset2cmt.iteritems():
-                idc.MakeComm(offset, '\n'.join(value))
-
-            return self.eval_expr(self.ir_arch.IRDst)
-
-
-    class Engine(SymbExecCTypeFix):
-        def __init__(self, state):
-            mychandler = MyCHandler(types_mngr, state.infos_types)
-            super(Engine, self).__init__(ir_arch,
-                                         state.symbols,
-                                         state.infos_types,
-                                         mychandler)
-
+            add_state(ir_arch, todo, states, son.offset,
+                      symbexec_engine.get_state())
 
     for lbl, state in states.iteritems():
-        symbexec_engine = Engine(state)
+        symbexec_engine = CTypeEngineFixer(ir_arch, state)
         addr = symbexec_engine.emul_ir_block(lbl)
         symbexec_engine.del_mem_above_stack(ir_arch.sp)
 
 
-
-
-
-
-
-
+if __name__ == "__main__":
+    analyse_function()