Merge pull request #147 from commial/printable_sem

Printable sem

4 files changed, 306 insertions, 147 deletions

diff --git a/miasm2/arch/mips32/sem.py b/miasm2/arch/mips32/sem.py
index 909b9b6c..0d62cc58 100644
--- a/miasm2/arch/mips32/sem.py
+++ b/miasm2/arch/mips32/sem.py
@@ -2,155 +2,133 @@ import miasm2.expression.expression as m2_expr
 from miasm2.ir.ir import ir, irbloc
 from miasm2.arch.mips32.arch import mn_mips32
 from miasm2.arch.mips32.regs import R_LO, R_HI, PC, RA
-
-
-def addiu(ir, instr, a, b, c):
-    """Adds a register @b and a sign-extended immediate value @c and stores the
-    result in a register @a"""
-    e = []
-    e.append(m2_expr.ExprAff(a, b+c))
-    return e, []
-
-def lw(ir, instr, a, b):
-    "A word is loaded into a register @a from the specified address @b."
-    e = []
-    e.append(m2_expr.ExprAff(a, b))
-    return e, []
-
-def sw(ir, instr, a, b):
-    "The contents of @b is stored at the specified address @a."
-    e = []
-    e.append(m2_expr.ExprAff(b, a))
-    return e, []
-
-def jal(ir, instr, a):
-    "Jumps to the calculated address @a and stores the return address in $RA"
-    e = []
-    n = m2_expr.ExprId(ir.get_next_break_label(instr))
-    e.append(m2_expr.ExprAff(PC, a))
-    e.append(m2_expr.ExprAff(ir.IRDst, a))
-    e.append(m2_expr.ExprAff(RA, n))
-    return e, []
-
-def jalr(ir, instr, a, b):
-    """Jump to an address stored in a register @a, and store the return address
-    in another register @b"""
-    e = []
-    n = m2_expr.ExprId(ir.get_next_break_label(instr))
-    e.append(m2_expr.ExprAff(PC, a))
-    e.append(m2_expr.ExprAff(ir.IRDst, a))
-    e.append(m2_expr.ExprAff(b, n))
-    return e, []
-
-def bal(ir, instr, a):
-    e = []
-    n = m2_expr.ExprId(ir.get_next_break_label(instr))
-    e.append(m2_expr.ExprAff(PC, a))
-    e.append(m2_expr.ExprAff(ir.IRDst, a))
-    e.append(m2_expr.ExprAff(RA, n))
-    return e, []
-
-def l_b(ir, instr, a):
-    e = []
-    e.append(m2_expr.ExprAff(PC, a))
-    e.append(m2_expr.ExprAff(ir.IRDst, a))
-    return e, []
-
-def lbu(ir, instr, a, b):
-    """A byte is loaded (unsigned extended) into a register @a from the
-    specified address @b."""
-    e = []
-    b = m2_expr.ExprMem(b.arg, 8)
-    e.append(m2_expr.ExprAff(a, b.zeroExtend(32)))
-    return e, []
-
-def lhu(ir, instr, a, b):
-    """A word is loaded (unsigned extended) into a register @a from the
-    specified address @b."""
-    e = []
-    b = m2_expr.ExprMem(b.arg, 16)
-    e.append(m2_expr.ExprAff(a, b.zeroExtend(32)))
-    return e, []
-
-
-def lb(ir, instr, a, b):
-    "A byte is loaded into a register @a from the specified address @b."
-    e = []
-    b = m2_expr.ExprMem(b.arg, 8)
-    e.append(m2_expr.ExprAff(a, b.signExtend(32)))
-    return e, []
-
-def beq(ir, instr, a, b, c):
-    "Branches on @c if the quantities of two registers @a, @b are equal"
-    e = []
-    n = m2_expr.ExprId(ir.get_next_break_label(instr))
-    dst_o = m2_expr.ExprCond(a-b, n, c)
-    e = [m2_expr.ExprAff(PC, dst_o),
-         m2_expr.ExprAff(ir.IRDst, dst_o)
-     ]
-    return e, []
-
-def bgez(ir, instr, a, b):
-    """Branches on @b if the quantities of register @a is greater than or equal
-    to zero"""
-    e = []
-    n = m2_expr.ExprId(ir.get_next_break_label(instr))
-    dst_o = m2_expr.ExprCond(a.msb(), n, b)
-    e = [m2_expr.ExprAff(PC, dst_o),
-         m2_expr.ExprAff(ir.IRDst, dst_o)
-     ]
-    return e, []
-
-def bne(ir, instr, a, b, c):
-    "Branches on @c if the quantities of two registers @a, @b are NOT equal"
-    e = []
-    n = m2_expr.ExprId(ir.get_next_break_label(instr))
-    dst_o = m2_expr.ExprCond(a-b, c, n)
-    e = [m2_expr.ExprAff(PC, dst_o),
-         m2_expr.ExprAff(ir.IRDst, dst_o)
-    ]
-    return e, []
-
-def lui(ir, instr, a, b):
-    """The immediate value @b is shifted left 16 bits and stored in the register
-    @a. The lower 16 bits are zeroes."""
-    e = []
-    e.append(m2_expr.ExprAff(a,
-                             m2_expr.ExprCompose([(m2_expr.ExprInt16(0), 0, 16),
-                                                  (b[:16], 16, 32)])))
-    return e, []
-
-def nop(ir, instr):
+from miasm2.core.sembuilder import SemBuilder
+
+
+# SemBuilder context
+ctx = {"R_LO": R_LO,
+       "R_HI": R_HI,
+       "PC": PC,
+       "RA": RA}
+ctx.update(m2_expr.__dict__)
+sbuild = SemBuilder(ctx)
+
+
+@sbuild.parse
+def addiu(Arg1, Arg2, Arg3):
+    """Adds a register @Arg3 and a sign-extended immediate value @Arg2 and
+    stores the result in a register @Arg1"""
+    Arg1 = Arg2 + Arg3
+
+@sbuild.parse
+def lw(Arg1, Arg2):
+    "A word is loaded into a register @Arg1 from the specified address @Arg2."
+    Arg1 = Arg2
+
+@sbuild.parse
+def sw(Arg1, Arg2):
+    "The contents of @Arg2 is stored at the specified address @Arg1."
+    Arg2 = Arg1
+
+@sbuild.parse
+def jal(Arg1):
+    "Jumps to the calculated address @Arg1 and stores the return address in $RA"
+    PC = Arg1
+    ir.IRDst = Arg1
+    RA = ExprId(ir.get_next_break_label(instr))
+
+@sbuild.parse
+def jalr(Arg1, Arg2):
+    """Jump to an address stored in a register @Arg1, and store the return
+    address in another register @Arg2"""
+    PC = Arg1
+    ir.IRDst = Arg1
+    Arg2 = ExprId(ir.get_next_break_label(instr))
+
+@sbuild.parse
+def bal(Arg1):
+    PC = Arg1
+    ir.IRDst = Arg1
+    RA = ExprId(ir.get_next_break_label(instr))
+
+@sbuild.parse
+def l_b(Arg1):
+    PC = Arg1
+    ir.IRDst = Arg1
+
+@sbuild.parse
+def lbu(Arg1, Arg2):
+    """A byte is loaded (unsigned extended) into a register @Arg1 from the
+    specified address @Arg2."""
+    Arg1 = mem8[Arg2.arg].zeroExtend(32)
+
+@sbuild.parse
+def lhu(Arg1, Arg2):
+    """A word is loaded (unsigned extended) into a register @Arg1 from the
+    specified address @Arg2."""
+    Arg1 = mem16[Arg2.arg].zeroExtend(32)
+
+@sbuild.parse
+def lb(Arg1, Arg2):
+    "A byte is loaded into a register @Arg1 from the specified address @Arg2."
+    Arg1 = mem8[Arg2.arg].signExtend(32)
+
+@sbuild.parse
+def beq(Arg1, Arg2, Arg3):
+    "Branches on @Arg3 if the quantities of two registers @Arg1, @Arg2 are eq"
+    dst = ExprId(ir.get_next_break_label(instr)) if Arg1 - Arg2 else Arg3
+    PC = dst
+    ir.IRDst = dst
+
+@sbuild.parse
+def bgez(Arg1, Arg2):
+    """Branches on @Arg2 if the quantities of register @Arg1 is greater than or
+    equal to zero"""
+    dst = ExprId(ir.get_next_break_label(instr)) if Arg1.msb() else Arg2
+    PC = dst
+    ir.IRDst = dst
+
+@sbuild.parse
+def bne(Arg1, Arg2, Arg3):
+    """Branches on @Arg3 if the quantities of two registers @Arg1, @Arg2 are NOT
+    equal"""
+    dst = Arg3 if Arg1 - Arg2 else ExprId(ir.get_next_break_label(instr))
+    PC = dst
+    ir.IRDst = dst
+
+@sbuild.parse
+def lui(Arg1, Arg2):
+    """The immediate value @Arg2 is shifted left 16 bits and stored in the
+    register @Arg1. The lower 16 bits are zeroes."""
+    Arg1 = ExprCompose([(i16(0), 0, 16), (Arg2[:16], 16, 32)])
+
+@sbuild.parse
+def nop():
     """Do nothing"""
-    return [], []
 
-def j(ir, instr, a):
-    """Jump to an address @a"""
-    e = []
-    e.append(m2_expr.ExprAff(PC, a))
-    e.append(m2_expr.ExprAff(ir.IRDst, a))
-    return e, []
-
-def l_or(ir, instr, a, b, c):
-    """Bitwise logical ors two registers @b, @c and stores the result in a
-    register @a"""
-    e = []
-    e.append(m2_expr.ExprAff(a, b|c))
-    return e, []
-
-def nor(ir, instr, a, b, c):
-    """Bitwise logical Nors two registers @b, @c and stores the result in a
-    register @a"""
-    e = []
-    e.append(m2_expr.ExprAff(a, (b|c)^m2_expr.ExprInt32(0xFFFFFFFF)))
-    return e, []
-
-def l_and(ir, instr, a, b, c):
-    """Bitwise logical ands two registers @b, @c and stores the result in a
-    register @a"""
-    e = []
-    e.append(m2_expr.ExprAff(a, b&c))
-    return e, []
+@sbuild.parse
+def j(Arg1):
+    """Jump to an address @Arg1"""
+    PC = Arg1
+    ir.IRDst = Arg1
+
+@sbuild.parse
+def l_or(Arg1, Arg2, Arg3):
+    """Bitwise logical ors two registers @Arg2, @Arg3 and stores the result in a
+    register @Arg1"""
+    Arg1 = Arg2 | Arg3
+
+@sbuild.parse
+def nor(Arg1, Arg2, Arg3):
+    """Bitwise logical Nors two registers @Arg2, @Arg3 and stores the result in
+    a register @Arg1"""
+    Arg1 = (Arg2 | Arg3) ^ i32(-1)
+
+@sbuild.parse
+def l_and(Arg1, Arg2, Arg3):
+    """Bitwise logical ands two registers @Arg2, @Arg3 and stores the result in
+    a register @Arg1"""
+    Arg1 = Arg2 & Arg3
 
 def ext(ir, instr, a, b, c, d):
     e = []
diff --git a/miasm2/core/sembuilder.py b/miasm2/core/sembuilder.py
new file mode 100644
index 00000000..34271960
--- /dev/null
+++ b/miasm2/core/sembuilder.py
@@ -0,0 +1,155 @@
+"Helper to quickly build instruction's semantic side effects"
+
+import inspect
+import ast
+import re
+
+
+class MiasmTransformer(ast.NodeTransformer):
+    """AST visitor translating DSL to Miasm expression
+
+    memX[Y]       -> ExprMem(Y, X)
+    iX(Y)         -> ExprIntX(Y)
+    X if Y else Z -> ExprCond(Y, X, Z)
+    """
+
+    # Parsers
+    parse_integer = re.compile("^i([0-9]+)$")
+    parse_mem = re.compile("^mem([0-9]+)$")
+
+    # Visitors
+
+    def visit_Call(self, node):
+        """iX(Y) -> ExprIntX(Y)"""
+        # Match the function name
+        if not isinstance(node.func, ast.Name):
+            # TODO: launch visitor on node
+            return node
+
+        fc_name = node.func.id
+        new_name = fc_name
+        integer = self.parse_integer.search(fc_name)
+
+        # Do replacement
+        if integer is not None:
+            new_name = "ExprInt%s" % integer.groups()[0]
+
+        # Replace in the node
+        node.func.id = new_name
+        # TODO: launch visitor on node
+        return node
+
+    def visit_Subscript(self, node):
+        """memX[Y] -> ExprMem(Y, X)"""
+
+        # Detect the syntax
+        if not isinstance(node.value, ast.Name):
+            return node
+        name = node.value.id
+        mem = self.parse_mem.search(name)
+        if mem is None:
+            # TODO: launch visitor on node
+            return node
+
+        # Do replacement
+        addr = self.visit(node.slice.value)
+        call = ast.Call(func=ast.Name(id='ExprMem', ctx=ast.Load()),
+                        args=[addr, ast.Num(n=int(mem.groups()[0]))],
+                        keywords=[], starargs=None, kwargs=None)
+        return call
+
+    def visit_IfExp(self, node):
+        """X if Y else Z -> ExprCond(Y, X, Z)"""
+        call = ast.Call(func=ast.Name(id='ExprCond', ctx=ast.Load()),
+                        args=[self.visit(node.test),
+                              self.visit(node.body),
+                              self.visit(node.orelse)],
+                        keywords=[], starargs=None, kwargs=None)
+        return call
+
+
+class SemBuilder(object):
+    """Helper for building instruction's semantic side effects method
+
+    This class provides a decorator @parse to use on them.
+    The context in which the function will be parsed must be supplied on
+    instanciation
+    """
+
+    def __init__(self, ctx):
+        """Create a SemBuilder
+        @ctx: context dictionnary used during parsing
+        """
+        self.ctx = dict(ctx)
+        self.transformer = MiasmTransformer()
+
+    def parse(self, func):
+        """Function decorator, returning a correct method from a pseudo-Python
+        one"""
+
+        # Get the function AST
+        parsed = ast.parse(inspect.getsource(func))
+        fc_ast = parsed.body[0]
+        argument_names = [name.id for name in fc_ast.args.args]
+        body = []
+
+        # AffBlock of the current instruction
+        new_ast = []
+        for statement in fc_ast.body:
+
+            if isinstance(statement, ast.Assign):
+                src = self.transformer.visit(statement.value)
+                dst = self.transformer.visit(statement.targets[0])
+
+                if (isinstance(dst, ast.Name) and
+                    dst.id not in argument_names and
+                    dst.id not in self.ctx):
+                    # Real variable declaration
+                    statement.value = src
+                    body.append(statement)
+                    continue
+
+                dst.ctx = ast.Load()
+
+                res = ast.Call(func=ast.Name(id='ExprAff',
+                                             ctx=ast.Load()),
+                               args=[dst, src],
+                               keywords=[],
+                               starargs=None,
+                               kwargs=None)
+
+                new_ast.append(res)
+
+            elif (isinstance(statement, ast.Expr) and
+                  isinstance(statement.value, ast.Str)):
+                # String (docstring, comment, ...) -> keep it
+                body.append(statement)
+            else:
+                # TODO: real var, +=, /=, -=, <<=, >>=, if/else, ...
+                raise RuntimeError("Unimplemented %s" % statement)
+
+        # Build the new function
+        fc_ast.args.args[0:0] = [ast.Name(id='ir', ctx=ast.Param()),
+                                 ast.Name(id='instr', ctx=ast.Param())]
+        body.append(ast.Return(value=ast.Tuple(elts=[ast.List(elts=new_ast,
+                                                              ctx=ast.Load()),
+                                                     ast.List(elts=[],
+                                                              ctx=ast.Load())],
+                                               ctx=ast.Load())))
+
+        ret = ast.Module([ast.FunctionDef(name=fc_ast.name,
+                                          args=fc_ast.args,
+                                          body=body,
+                                          decorator_list=[])])
+
+        # To display the generated function, use codegen.to_source
+        # codegen: https://github.com/andreif/codegen
+
+        # Compile according to the context
+        fixed = ast.fix_missing_locations(ret)
+        codeobj = compile(fixed, '<string>', 'exec')
+        ctx = self.ctx.copy()
+        eval(codeobj, ctx)
+
+        # Get the function back
+        return ctx[fc_ast.name]
diff --git a/test/core/sembuilder.py b/test/core/sembuilder.py
new file mode 100644
index 00000000..3c1b7d3f
--- /dev/null
+++ b/test/core/sembuilder.py
@@ -0,0 +1,25 @@
+import inspect
+from pdb import pm
+
+from miasm2.core.sembuilder import SemBuilder
+import miasm2.expression.expression as m2_expr
+
+sb = SemBuilder(m2_expr.__dict__)
+
+@sb.parse
+def test(Arg1, Arg2, Arg3):
+    "Test docstring"
+    Arg1 = Arg2
+    mem32[Arg1] = Arg2
+    mem32[Arg2] = Arg3  + i32(4) - mem32[Arg1]
+    Arg3 = Arg3 if Arg2 else i32(0)
+    tmpvar = i32(2)
+    Arg2 = tmpvar
+
+a = m2_expr.ExprId('A')
+b = m2_expr.ExprId('B')
+c = m2_expr.ExprId('C')
+ir = None
+instr = None
+print test(ir, instr, a, b, c)
+print test.__doc__
diff --git a/test/test_all.py b/test/test_all.py
index 52f2238b..aefd0196 100644
--- a/test/test_all.py
+++ b/test/test_all.py
@@ -101,6 +101,7 @@ for script in ["interval.py",
                "graph.py",
                "parse_asm.py",
                "utils.py",
+               "sembuilder.py",
                ]:
     testset += RegressionTest([script], base_dir="core")
 ## Expression