Merge pull request #596 from serpilliere/fix_ctype

Fix ctype

7 files changed, 199 insertions, 133 deletions

diff --git a/example/expression/access_c.py b/example/expression/access_c.py
index 7255e23a..f285eb55 100644
--- a/example/expression/access_c.py
+++ b/example/expression/access_c.py
@@ -1,16 +1,3 @@
-import sys
-
-from miasm2.analysis.machine import Machine
-from miasm2.analysis.binary import Container
-from miasm2.expression.expression import ExprOp, ExprCompose, ExprId
-from miasm2.analysis.depgraph import DependencyGraph
-
-from miasm2.arch.x86.ctype import CTypeAMD64_unk
-
-from miasm2.core.objc import CTypeAnalyzer, ExprToAccessC, CHandler
-from miasm2.core.objc import CTypesManagerNotPacked
-from miasm2.core.ctypesmngr import CAstTypes, CTypePtr, CTypeStruct
-
 """
 
 This example demonstrates the recovering of possible C types for an arbitrary
@@ -54,6 +41,20 @@ ExprCompose(var1, 0) => var1
 """
 
 
+import sys
+
+from miasm2.analysis.machine import Machine
+from miasm2.analysis.binary import Container
+from miasm2.expression.expression import ExprOp, ExprCompose, ExprId, ExprInt
+from miasm2.analysis.depgraph import DependencyGraph
+
+from miasm2.arch.x86.ctype import CTypeAMD64_unk
+
+from miasm2.core.objc import ExprToAccessC, CHandler
+from miasm2.core.objc import CTypesManagerNotPacked
+from miasm2.core.ctypesmngr import CAstTypes, CTypePtr, CTypeStruct
+
+
 def find_call(ira):
     """Returns (irb, index) which call"""
 
@@ -74,24 +75,10 @@ def find_call(ira):
         yield irb, index
 
 
-class MyCTypeAnalyzer(CTypeAnalyzer):
-    """Custom CTypeAnalyzer to complete type analysis"""
-
-    def reduce_compose(self, node, _):
-        """Custom reduction rule: {XXX, 0} -> typeof(XXX)"""
-        if not (isinstance(node.expr, ExprCompose) and
-                len(node.expr.args) == 2 and
-                node.expr.args[1].is_int(0)):
-            return None
-        return node.args[0].info
-
-    reduction_rules = CTypeAnalyzer.reduction_rules + [reduce_compose]
-
-
 class MyExprToAccessC(ExprToAccessC):
     """Custom ExprToAccessC to complete expression traduction to C"""
 
-    def reduce_compose(self, node, _):
+    def reduce_compose(self, node, **kwargs):
         """Custom reduction rule: {XXX, 0} -> XXX"""
         if not (isinstance(node.expr, ExprCompose) and
                 len(node.expr.args) == 2 and
@@ -123,7 +110,6 @@ def get_funcs_arg0(ctx, ira, lbl_head):
 class MyCHandler(CHandler):
     """Custom CHandler to add complementary C handling rules"""
 
-    cTypeAnalyzer_cls = MyCTypeAnalyzer
     exprToAccessC_cls = MyExprToAccessC
 
 
@@ -170,17 +156,13 @@ open('graph_irflow.dot', 'w').write(ir_arch_a.graph.dot())
 ptr_llhuman = types_mngr.get_objc(CTypePtr(CTypeStruct('ll_human')))
 arg0 = ExprId('ptr', 64)
 ctx = {ir_arch_a.arch.regs.RDI: arg0}
-expr_types = {arg0.name: ptr_llhuman}
+expr_types = {arg0: (ptr_llhuman,),
+              ExprInt(0x8A, 64): (ptr_llhuman,)}
 
 mychandler = MyCHandler(types_mngr, expr_types)
 
 for expr in get_funcs_arg0(ctx, ir_arch_a, lbl_head):
     print "Access:", expr
-    target_types = mychandler.expr_to_types(expr)
-    for target_type in target_types:
-        print '\tType:', target_type
-    c_strs = mychandler.expr_to_c(expr)
-    for c_str in c_strs:
-        print "\tC access:", c_str
-    print
-
+    for c_str, ctype in mychandler.expr_to_c_and_types(expr):
+        print '\taccess:', c_str
+        print '\tc type:', ctype
diff --git a/example/expression/constant_propagation.py b/example/expression/constant_propagation.py
new file mode 100644
index 00000000..70394580
--- /dev/null
+++ b/example/expression/constant_propagation.py
@@ -0,0 +1,54 @@
+"""
+Example of "constant expression" propagation.
+A "constant expression" is an expression based on constants or init regs.
+
+"""
+
+from argparse import ArgumentParser
+
+from miasm2.arch.x86.disasm import dis_x86_32 as dis_engine
+from miasm2.analysis.machine import Machine
+from miasm2.analysis.binary import Container
+from miasm2.analysis.cst_propag import propagate_cst_expr
+from miasm2.analysis.data_flow import dead_simp
+from miasm2.expression.simplifications import expr_simp
+
+
+parser = ArgumentParser("Constant expression propagation")
+parser.add_argument('filename', help="File to analyze")
+parser.add_argument('address', help="Starting address for disassembly engine")
+parser.add_argument('-s', "--simplify", action="store_true",
+                    help="Apply simplifications rules (liveness, graph simplification, ...)")
+
+args = parser.parse_args()
+
+
+machine = Machine("x86_32")
+
+cont = Container.from_stream(open(args.filename))
+ira, dis_engine = machine.ira, machine.dis_engine
+mdis = dis_engine(cont.bin_stream)
+ir_arch = ira(mdis.symbol_pool)
+addr = int(args.address, 0)
+
+
+blocks = mdis.dis_multiblock(addr)
+for block in blocks:
+    ir_arch.add_block(block)
+
+
+init_infos = ir_arch.arch.regs.regs_init
+cst_propag_link = propagate_cst_expr(ir_arch, addr, init_infos)
+
+if args.simplify:
+    ir_arch.simplify(expr_simp)
+    modified = True
+    while modified:
+        modified = False
+        modified |= dead_simp(ir_arch)
+        modified |= ir_arch.remove_empty_assignblks()
+        modified |= ir_arch.remove_jmp_blocks()
+        modified |= ir_arch.merge_blocks()
+
+
+open("%s.propag.dot" % args.filename, 'w').write(ir_arch.graph.dot())
diff --git a/example/expression/expr_c.py b/example/expression/expr_c.py
index 7adc7b50..ca92153a 100644
--- a/example/expression/expr_c.py
+++ b/example/expression/expr_c.py
@@ -41,9 +41,8 @@ types_mngr = CTypesManagerNotPacked(types_ast, base_types)
 ptr_rectangle = types_mngr.get_objc(CTypePtr(CTypeStruct('rectangle')))
 
 ptr = ExprId('ptr', 64)
-expr_types = {ptr.name: ptr_rectangle}
-
-mychandler = CHandler(types_mngr, expr_types)
+c_context = {ptr.name: ptr_rectangle}
+mychandler = CHandler(types_mngr, {})
 
 # Parse some C accesses
 c_acceses = ["ptr->width",
@@ -55,8 +54,8 @@ c_acceses = ["ptr->width",
             ]
 
 for c_str in c_acceses:
-    expr = mychandler.c_to_expr(c_str)
-    c_type = mychandler.c_to_type(c_str)
+    expr = mychandler.c_to_expr(c_str, c_context)
+    c_type = mychandler.c_to_type(c_str, c_context)
     print 'C access:', c_str
     print '\tExpr:', expr
     print '\tType:', c_type
diff --git a/example/expression/expr_reduce.py b/example/expression/expr_reduce.py
index b5fc96c8..bb94ceb9 100644
--- a/example/expression/expr_reduce.py
+++ b/example/expression/expr_reduce.py
@@ -20,7 +20,7 @@ class StructLookup(ExprReducer):
     FIELD_A_PTR = "FIELD_A_PTR"
     FIELD_A = "FIELD_A"
 
-    def reduce_int(self, node, _):
+    def reduce_int(self, node, **kwargs):
         """
         Reduction: int -> CST
         """
@@ -28,7 +28,7 @@ class StructLookup(ExprReducer):
             return self.CST
         return None
 
-    def reduce_ptr_struct(self, node, _):
+    def reduce_ptr_struct(self, node, **kwargs):
         """
         Reduction: ECX -> FIELD_A_PTR
         """
@@ -36,7 +36,7 @@ class StructLookup(ExprReducer):
             return self.FIELD_A_PTR
         return None
 
-    def reduce_ptr_plus_int(self, node, _):
+    def reduce_ptr_plus_int(self, node, **kwargs):
         """
         Reduction: ECX + CST -> FIELD_A_PTR
         """
@@ -46,7 +46,7 @@ class StructLookup(ExprReducer):
             return self.FIELD_A_PTR
         return None
 
-    def reduce_cst_op(self, node, _):
+    def reduce_cst_op(self, node, **kwargs):
         """
         Reduction: CST + CST -> CST
         """
@@ -56,7 +56,7 @@ class StructLookup(ExprReducer):
             return self.CST
         return None
 
-    def reduce_at_struct_ptr(self, node, _):
+    def reduce_at_struct_ptr(self, node, **kwargs):
         """
         Reduction: @FIELD_A_PTR -> FIELD_A
         """
diff --git a/example/expression/solve_condition_stp.py b/example/expression/solve_condition_stp.py
index c9d4c7af..438188ab 100644
--- a/example/expression/solve_condition_stp.py
+++ b/example/expression/solve_condition_stp.py
@@ -11,7 +11,7 @@ from miasm2.core.bin_stream import bin_stream_str
 from miasm2.core import asmblock
 from miasm2.expression.expression import get_rw
 from miasm2.expression.modint import uint32
-from miasm2.ir.symbexec import SymbolicExecutionEngine
+from miasm2.ir.symbexec import SymbolicExecutionEngine, get_block
 from miasm2.expression.simplifications import expr_simp
 from miasm2.expression import stp
 from miasm2.core import parse_asm
@@ -30,21 +30,6 @@ if not args:
     sys.exit(0)
 
 
-def get_block(ir_arch, mdis, ad):
-    if isinstance(ad, asmblock.AsmLabel):
-        l = ad
-    else:
-        l = mdis.symbol_pool.getby_offset_create(ad)
-    if not l in ir_arch.blocks:
-        ad = l.offset
-        b = mdis.dis_block(ad)
-        ir_arch.add_block(b)
-    b = ir_arch.get_block(l)
-    if b is None:
-        raise LookupError('no block found at that address: %s' % l)
-    return b
-
-
 def emul_symb(ir_arch, mdis, states_todo, states_done):
     while states_todo:
         ad, symbols, conds = states_todo.pop()
diff --git a/example/ida/ctype_propagation.py b/example/ida/ctype_propagation.py
index cb342213..7eb209cd 100644
--- a/example/ida/ctype_propagation.py
+++ b/example/ida/ctype_propagation.py
@@ -10,13 +10,15 @@ from miasm2.expression import expression as m2_expr
 from miasm2.expression.simplifications import expr_simp
 from miasm2.analysis.depgraph import DependencyGraph
 from miasm2.ir.ir import IRBlock, AssignBlock
-from miasm2.arch.x86.ctype import CTypeAMD64_unk
+from miasm2.arch.x86.ctype import CTypeAMD64_unk, CTypeX86_unk
 from miasm2.expression.expression import ExprId
-from miasm2.core.objc import CTypesManagerNotPacked, CTypeAnalyzer, ExprToAccessC, CHandler
+from miasm2.core.objc import CTypesManagerNotPacked, ExprToAccessC, CHandler
 from miasm2.core.ctypesmngr import CAstTypes
 from miasm2.expression.expression import ExprMem, ExprId, ExprInt, ExprOp, ExprAff
 from miasm2.ir.symbexec_types import SymbExecCType
 from miasm2.expression.parser import str_to_expr
+from miasm2.ir.symbexec import SymbolicExecutionEngine, SymbolicState
+from miasm2.analysis.cst_propag import add_state, propagate_cst_expr
 
 from utils import guess_machine
 
@@ -25,7 +27,6 @@ class TypePropagationForm(ida_kernwin.Form):
     def __init__(self, ira):
 
         self.ira = ira
-        self.stk_unalias_force = False
 
         default_types_info = r"""ExprId("RDX", 64): char *"""
         archs = ["AMD64_unk", "X86_32_unk"]
@@ -35,41 +36,31 @@ class TypePropagationForm(ida_kernwin.Form):
 BUTTON CANCEL NONE
 Dependency Graph Settings
 <##Header file :{headerFile}>
-<Architecture/complator:{cbReg}>
+<Architecture/complator:{arch}>
 <Types informations:{strTypesInfo}>
-<Unalias stack:{rUnaliasStack}>{cMethod}>
+<Unalias stack:{rUnaliasStack}>{cUnalias}>
 """, {
                           'headerFile': ida_kernwin.Form.FileInput(swidth=20, open=True),
-                          'cbReg': ida_kernwin.Form.DropdownListControl(
+                          'arch': ida_kernwin.Form.DropdownListControl(
                               items=archs,
                               readonly=False,
                               selval=archs[0]),
                           'strTypesInfo': ida_kernwin.Form.MultiLineTextControl(text=default_types_info,
                                                                     flags=ida_kernwin.Form.MultiLineTextControl.TXTF_FIXEDFONT),
-                          'cMethod': ida_kernwin.Form.ChkGroupControl(("rUnaliasStack",)),
+                          'cUnalias': ida_kernwin.Form.ChkGroupControl(("rUnaliasStack",)),
                       })
-        self.Compile()
+        form, args = self.Compile()
+        form.rUnaliasStack.checked = True
 
-    @property
-    def unalias_stack(self):
-        return self.cMethod.value & 1 or self.stk_unalias_force
 
-
-def get_block(ir_arch, mdis, addr):
-    """Get IRBlock at address @addr"""
-    lbl = ir_arch.get_label(addr)
-    if not lbl in ir_arch.blocks:
-        block = mdis.dis_block(lbl.offset)
-        ir_arch.add_block(block)
-    irblock = ir_arch.get_block(lbl)
-    if irblock is None:
-        raise LookupError('No block found at that address: %s' % lbl)
-    return irblock
-
-
-def get_types_mngr(headerFile):
+def get_types_mngr(headerFile, arch):
     text = open(headerFile).read()
-    base_types = CTypeAMD64_unk()
+    if arch == "AMD64_unk":
+        base_types = CTypeAMD64_unk()
+    elif arch =="X86_32_unk":
+        base_types = CTypeX86_unk()
+    else:
+        raise NotImplementedError("Unsupported arch")
     types_ast = CAstTypes()
 
     # Add C types definition
@@ -79,16 +70,11 @@ def get_types_mngr(headerFile):
     return types_mngr
 
 
-class MyCTypeAnalyzer(CTypeAnalyzer):
-    allow_none_result = True
-
-
 class MyExprToAccessC(ExprToAccessC):
     allow_none_result = True
 
 
 class MyCHandler(CHandler):
-    cTypeAnalyzer_cls = MyCTypeAnalyzer
     exprToAccessC_cls = MyExprToAccessC
 
 
@@ -103,52 +89,89 @@ class TypePropagationEngine(SymbExecCType):
 
 class SymbExecCTypeFix(SymbExecCType):
 
+    def __init__(self, ir_arch,
+                 symbols, chandler,
+                 cst_propag_link,
+                 func_read=None, func_write=None,
+                 sb_expr_simp=expr_simp):
+        super(SymbExecCTypeFix, self).__init__(ir_arch,
+                                               symbols,
+                                               chandler,
+                                               func_read=func_read,
+                                               func_write=func_write,
+                                               sb_expr_simp=expr_simp)
+
+        self.cst_propag_link = cst_propag_link
+
     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:
+        for index, assignblk in enumerate(irb.irs):
+            if set(assignblk) == set([self.ir_arch.IRDst, self.ir_arch.pc]):
+                # Don't display on jxx
+                continue
             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
+            tmp_r = assignblk.get_r()
+            tmp_w = assignblk.get_w()
+
+            todo = set()
+
+            # Replace PC with value to match IR args
+            pc_fixed = {self.ir_arch.pc: m2_expr.ExprInt(instr.offset + instr.l, self.ir_arch.pc.size)}
+            for arg in tmp_r:
+                arg = expr_simp(arg.replace_expr(pc_fixed))
+                if arg in tmp_w and not arg.is_mem():
+                    continue
+                todo.add(arg)
+
+            for expr in todo:
+                if expr.is_int():
+                    continue
+                for c_str, c_type in self.chandler.expr_to_c_and_types(expr, self.symbols):
+                    expr = self.cst_propag_link.get((irb.label, index), {}).get(expr, expr)
                     offset2cmt.setdefault(instr.offset, set()).add(
-                        "%s: %s" % (arg, str(objc)))
-            self.eval_ir(assignblk)
+                        "\n%s: %s\n%s" % (expr, c_str, c_type))
 
+            self.eval_ir(assignblk)
         for offset, value in offset2cmt.iteritems():
             idc.MakeComm(offset, '\n'.join(value))
+            print "%x\n" % offset, '\n'.join(value)
 
         return self.eval_expr(self.ir_arch.IRDst)
 
 
 class CTypeEngineFixer(SymbExecCTypeFix):
 
-    def __init__(self, ir_arch, types_mngr, state):
+    def __init__(self, ir_arch, types_mngr, state, cst_propag_link):
         mychandler = MyCHandler(types_mngr, state.symbols)
         super(CTypeEngineFixer, self).__init__(ir_arch,
                                                state.symbols,
-                                               mychandler)
+                                               mychandler,
+                                               cst_propag_link)
 
 
-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 get_ira_call_fixer(ira):
+
+    class iraCallStackFixer(ira):
+
+        def call_effects(self, ad, instr):
+            print hex(instr.offset), instr
+            stk_before = idc.GetSpd(instr.offset)
+            stk_after = idc.GetSpd(instr.offset + instr.l)
+            stk_diff = stk_after - stk_before
+            print hex(stk_diff)
+            return [AssignBlock([ExprAff(self.ret_reg, ExprOp('call_func_ret', ad)),
+                                 ExprAff(self.sp, self.sp + ExprInt(stk_diff, self.sp.size))
+                                 ],
+                                instr
+                                )]
+
+    return iraCallStackFixer
 
 
 def analyse_function():
@@ -159,7 +182,11 @@ def analyse_function():
 
     bs = bin_stream_ida()
     mdis = dis_engine(bs, dont_dis_nulstart_bloc=True)
-    ir_arch = ira(mdis.symbol_pool)
+
+
+    iraCallStackFixer = get_ira_call_fixer(ira)
+    ir_arch = iraCallStackFixer(mdis.symbol_pool)
+
 
     # Get the current function
     func = ida_funcs.get_func(idc.ScreenEA())
@@ -169,13 +196,20 @@ def analyse_function():
     for block in blocks:
         ir_arch.add_block(block)
 
+
     # Get settings
     settings = TypePropagationForm(ir_arch)
     ret = settings.Execute()
     if not ret:
         return
 
-    types_mngr = get_types_mngr(settings.headerFile.value)
+    cst_propag_link = {}
+    if settings.cUnalias.value:
+        init_infos = {ir_arch.sp: ir_arch.arch.regs.regs_init[ir_arch.sp] }
+        cst_propag_link = propagate_cst_expr(ir_arch, addr, init_infos)
+
+
+    types_mngr = get_types_mngr(settings.headerFile.value, settings.arch.value)
     mychandler = MyCHandler(types_mngr, {})
     infos_types = {}
     for line in settings.strTypesInfo.value.split('\n'):
@@ -184,14 +218,13 @@ def analyse_function():
         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(
+        ast = mychandler.types_mngr.types_ast.parse_c_type(
             ctype_str)
-        ctype = mychandler.type_analyzer.types_mngr.types_ast.ast_parse_declaration(ast.ext[
-                                                                                    0])
+        ctype = mychandler.types_mngr.types_ast.ast_parse_declaration(ast.ext[0])
         objc = types_mngr.get_objc(ctype)
         print '=' * 20
         print expr, objc
-        infos_types[expr] = objc
+        infos_types[expr] = set([objc])
 
     # Add fake head
     lbl_real_start = ir_arch.symbol_pool.getby_offset(addr)
@@ -220,21 +253,18 @@ def analyse_function():
         done.add((lbl, state))
         symbexec_engine = TypePropagationEngine(ir_arch, types_mngr, state)
 
-        get_block(ir_arch, mdis, lbl)
-
+        assert lbl in ir_arch.blocks
         addr = symbexec_engine.emul_ir_block(lbl)
         symbexec_engine.del_mem_above_stack(ir_arch.sp)
 
         ir_arch._graph = None
         sons = ir_arch.graph.successors(lbl)
         for son in sons:
-            if son.offset is None:
-                continue
-            add_state(ir_arch, todo, states, son.offset,
+            add_state(ir_arch, todo, states, son,
                       symbexec_engine.get_state())
 
     for lbl, state in states.iteritems():
-        symbexec_engine = CTypeEngineFixer(ir_arch, types_mngr, state)
+        symbexec_engine = CTypeEngineFixer(ir_arch, types_mngr, state, cst_propag_link)
         addr = symbexec_engine.emul_ir_block(lbl)
         symbexec_engine.del_mem_above_stack(ir_arch.sp)
 
diff --git a/example/samples/human.S b/example/samples/human.S
index 750aa5b7..6cdeab0f 100644
--- a/example/samples/human.S
+++ b/example/samples/human.S
@@ -1,11 +1,11 @@
 	;;  Walk a human link list and print its information
 main:
 	TEST       RDI, RDI
-	JZ         end
+	JZ         next
 
 	PUSH       RBX
 	MOV        RBX, RDI
-loop:
+loop_arg:
 	LEA        RSI, QWORD PTR [RBX+0x10]
 	LEA        RDI, QWORD PTR [name-_+RIP]
 	XOR        EAX, EAX
@@ -23,9 +23,23 @@ loop:
 
 	MOV        RBX, QWORD PTR [RBX]
 	TEST       RBX, RBX
-	JNZ        loop
+	JNZ        loop_arg
 
 	POP        RBX
+next:
+
+
+	LEA        RBX, QWORD PTR [struct_human_ptr-_+RIP]
+loop_global:
+	CMP        RBX, 0
+	JZ         end
+
+	LEA        RSI, QWORD PTR [RBX+0x10]
+	LEA        RDI, QWORD PTR [name-_+RIP]
+	XOR        EAX, EAX
+	CALL       printf
+	MOV        RBX, QWORD PTR [RBX]
+	JMP        loop_global
 end:
 	RET
 
@@ -39,3 +53,5 @@ height:
 .string "Height: %d\n"
 name:
 .string "Name: %s\n"
+struct_human_ptr:
+.dword 0xdead, 0xcafe