Merge pull request #245 from serpilliere/fix_parse_asm

Fix parse asm

5 files changed, 74 insertions, 5 deletions

diff --git a/test/arch/aarch64/unit/asm_test.py b/test/arch/aarch64/unit/asm_test.py
index 60ed418e..9e0d5ea8 100644
--- a/test/arch/aarch64/unit/asm_test.py
+++ b/test/arch/aarch64/unit/asm_test.py
@@ -41,7 +41,7 @@ class Asm_Test(object):
         # fix shellcode addr
         symbol_pool.set_offset(symbol_pool.getby_name("main"), 0x0)
         s = StrPatchwork()
-        patches = asmbloc.asm_resolve_final(mn_aarch64, blocs[0], symbol_pool)
+        patches = asmbloc.asm_resolve_final(mn_aarch64, blocs, symbol_pool)
         for offset, raw in patches.items():
             s[offset] = raw
 
diff --git a/test/arch/mips32/unit/asm_test.py b/test/arch/mips32/unit/asm_test.py
index b6cb7b2d..a00d0842 100644
--- a/test/arch/mips32/unit/asm_test.py
+++ b/test/arch/mips32/unit/asm_test.py
@@ -41,7 +41,7 @@ class Asm_Test(object):
         # fix shellcode addr
         symbol_pool.set_offset(symbol_pool.getby_name("main"), 0x0)
         s = StrPatchwork()
-        patches = asmbloc.asm_resolve_final(mn_mips32, blocs[0], symbol_pool)
+        patches = asmbloc.asm_resolve_final(mn_mips32, blocs, symbol_pool)
         for offset, raw in patches.items():
             s[offset] = raw
 
diff --git a/test/arch/x86/sem.py b/test/arch/x86/sem.py
index b80ab33d..617b929b 100644
--- a/test/arch/x86/sem.py
+++ b/test/arch/x86/sem.py
@@ -47,9 +47,9 @@ def compute(ir, mode, asm, inputstate={}, debug=False):
 def compute_txt(ir, mode, txt, inputstate={}, debug=False):
     blocs, symbol_pool = parse_asm.parse_txt(mn, mode, txt)
     symbol_pool.set_offset(symbol_pool.getby_name("main"), 0x0)
-    patches = asmbloc.asm_resolve_final(mn, blocs[0], symbol_pool)
+    patches = asmbloc.asm_resolve_final(mn, blocs, symbol_pool)
     interm = ir(symbol_pool)
-    for bbl in blocs[0]:
+    for bbl in blocs:
         interm.add_bloc(bbl)
     return symb_exec(interm, inputstate, debug)
 
diff --git a/test/arch/x86/unit/asm_test.py b/test/arch/x86/unit/asm_test.py
index c6381d9e..bf609aa5 100644
--- a/test/arch/x86/unit/asm_test.py
+++ b/test/arch/x86/unit/asm_test.py
@@ -41,7 +41,7 @@ class Asm_Test(object):
         # fix shellcode addr
         symbol_pool.set_offset(symbol_pool.getby_name("main"), 0x0)
         s = StrPatchwork()
-        patches = asmbloc.asm_resolve_final(mn_x86, blocs[0], symbol_pool)
+        patches = asmbloc.asm_resolve_final(mn_x86, blocs, symbol_pool)
         for offset, raw in patches.items():
             s[offset] = raw
 
diff --git a/test/core/parse_asm.py b/test/core/parse_asm.py
index c2a6dc72..a488d075 100644
--- a/test/core/parse_asm.py
+++ b/test/core/parse_asm.py
@@ -35,6 +35,75 @@ class TestParseAsm(unittest.TestCase):
         self.assertTrue(parse_txt(mn_x86, 32, ASM0))
         self.assertRaises(ValueError, parse_txt, mn_x86, 32, ASM1)
 
+    def test_DirectiveDontSplit(self):
+        from miasm2.arch.x86.arch import mn_x86
+        from miasm2.core.parse_asm import parse_txt
+        from miasm2.core.asmbloc import asm_resolve_final
+
+        ASM0 = '''
+        lbl0:
+            INC   EAX
+            JNZ   lbl0
+            INC   EAX
+            JZ    lbl2
+        lbl1:
+            NOP
+            JMP   lbl0
+        .dontsplit
+        lbl2:
+            MOV   EAX, ECX
+            RET
+        .dontsplit
+        lbl3:
+            ADD   EAX, EBX
+        .dontsplit
+        lbl4:
+        .align 0x10
+        .string "test"
+        lbl5:
+        .string "toto"
+        '''
+
+        blocks, symbol_pool = parse_txt(mn_x86, 32, ASM0)
+        patches = asm_resolve_final(mn_x86,
+                                    blocks,
+                                    symbol_pool)
+        lbls = []
+        for i in xrange(6):
+            lbls.append(symbol_pool.getby_name('lbl%d' % i))
+        # align test
+        assert(lbls[5].offset % 0x10 == 0)
+        lbl2block = {}
+        for block in blocks:
+            lbl2block[block.label] = block
+        # dontsplit test
+        assert(lbls[2] == lbl2block[lbls[1]].get_next())
+        assert(lbls[3] == lbl2block[lbls[2]].get_next())
+        assert(lbls[4] == lbl2block[lbls[3]].get_next())
+        assert(lbls[5] == lbl2block[lbls[4]].get_next())
+
+    def test_DirectiveSplit(self):
+        from miasm2.arch.x86.arch import mn_x86
+        from miasm2.core.parse_asm import parse_txt
+
+        ASM0 = '''
+        lbl0:
+            JNZ   lbl0
+        .split
+        lbl1:
+            RET
+        '''
+
+        blocks, symbol_pool = parse_txt(mn_x86, 32, ASM0)
+        lbls = []
+        for i in xrange(2):
+            lbls.append(symbol_pool.getby_name('lbl%d' % i))
+        lbl2block = {}
+        for block in blocks:
+            lbl2block[block.label] = block
+        # split test
+        assert(lbl2block[lbls[1]].get_next() is None)
+
 if __name__ == '__main__':
     testsuite = unittest.TestLoader().loadTestsFromTestCase(TestParseAsm)
     report = unittest.TextTestRunner(verbosity=2).run(testsuite)