diff options
Diffstat (limited to 'test/core')
| -rw-r--r-- | test/core/asmblock.py | 34 | ||||
| -rw-r--r-- | test/core/locationdb.py | 108 | ||||
| -rwxr-xr-x | test/core/parse_asm.py | 12 | ||||
| -rw-r--r-- | test/core/sembuilder.py | 10 |
4 files changed, 136 insertions, 28 deletions
diff --git a/test/core/asmblock.py b/test/core/asmblock.py index 0e965bfd..cd1d262a 100644 --- a/test/core/asmblock.py +++ b/test/core/asmblock.py @@ -105,7 +105,7 @@ open("graph2.dot", "w").write(asmcfg.dot()) # Test helper methods ## loc_key_to_block should always be updated assert asmcfg.loc_key_to_block(first_block.loc_key) == first_block -testlabel = mdis.symbol_pool.getby_name_create("testlabel") +testlabel = mdis.loc_db.get_or_create_name_location("testlabel") my_block = AsmBlock(testlabel) asmcfg.add_block(my_block) assert len(asmcfg) == 3 @@ -116,7 +116,7 @@ assert asmcfg.loc_key_to_block(my_block.loc_key) == my_block assert len(list(asmcfg.get_bad_blocks())) == 0 assert len(list(asmcfg.get_bad_blocks_predecessors())) == 0 ### Add a bad block, not linked -testlabel_bad = mdis.symbol_pool.getby_name_create("testlabel_bad") +testlabel_bad = mdis.loc_db.get_or_create_name_location("testlabel_bad") my_bad_block = AsmBlockBad(testlabel_bad) asmcfg.add_block(my_bad_block) assert list(asmcfg.get_bad_blocks()) == [my_bad_block] @@ -135,7 +135,7 @@ assert len(list(asmcfg.get_bad_blocks_predecessors(strict=True))) == 0 ## Sanity check asmcfg.sanity_check() ### Next on itself -testlabel_nextitself = mdis.symbol_pool.getby_name_create("testlabel_nextitself") +testlabel_nextitself = mdis.loc_db.get_or_create_name_location("testlabel_nextitself") my_block_ni = AsmBlock(testlabel_nextitself) my_block_ni.bto.add(AsmConstraintNext(my_block_ni.loc_key)) asmcfg.add_block(my_block_ni) @@ -149,11 +149,11 @@ assert error_raised asmcfg.del_block(my_block_ni) asmcfg.sanity_check() ### Multiple next on the same node -testlabel_target = mdis.symbol_pool.getby_name_create("testlabel_target") +testlabel_target = mdis.loc_db.get_or_create_name_location("testlabel_target") my_block_target = AsmBlock(testlabel_target) asmcfg.add_block(my_block_target) -testlabel_src1 = mdis.symbol_pool.getby_name_create("testlabel_src1") -testlabel_src2 = mdis.symbol_pool.getby_name_create("testlabel_src2") +testlabel_src1 = mdis.loc_db.get_or_create_name_location("testlabel_src1") +testlabel_src2 = mdis.loc_db.get_or_create_name_location("testlabel_src2") my_block_src1 = AsmBlock(testlabel_src1) my_block_src2 = AsmBlock(testlabel_src2) my_block_src1.bto.add(AsmConstraintNext(my_block_target.loc_key)) @@ -184,8 +184,8 @@ assert asmcfg.loc_key_to_block(my_block_src1.loc_key).max_size == 0 ## Check pendings ### Create a pending element -testlabel_pend_src = mdis.symbol_pool.getby_name_create("testlabel_pend_src") -testlabel_pend_dst = mdis.symbol_pool.getby_name_create("testlabel_pend_dst") +testlabel_pend_src = mdis.loc_db.get_or_create_name_location("testlabel_pend_src") +testlabel_pend_dst = mdis.loc_db.get_or_create_name_location("testlabel_pend_dst") my_block_src = AsmBlock(testlabel_pend_src) my_block_dst = AsmBlock(testlabel_pend_dst) my_block_src.bto.add(AsmConstraintTo(my_block_dst.loc_key)) @@ -247,7 +247,7 @@ assert len(entry_block.lines) == 4 assert map(str, entry_block.lines) == ['XOR EAX, EAX', 'XOR EBX, EBX', 'XOR ECX, ECX', - 'JNZ loc_3'] + 'JNZ loc_key_3'] assert len(asmcfg.successors(entry_block.loc_key)) == 2 assert len(entry_block.bto) == 2 nextb = asmcfg.loc_key_to_block((cons.loc_key for cons in entry_block.bto @@ -258,26 +258,26 @@ assert len(nextb.lines) == 4 assert map(str, nextb.lines) == ['XOR EDX, EDX', 'XOR ESI, ESI', 'XOR EDI, EDI', - 'JMP loc_4'] + 'JMP loc_key_4'] assert asmcfg.successors(nextb.loc_key) == [nextb.loc_key] assert len(tob.lines) == 2 assert map(str, tob.lines) == ['XOR EBP, EBP', - 'JMP loc_3'] + 'JMP loc_key_3'] assert asmcfg.successors(tob.loc_key) == [tob.loc_key] # Check split_block ## Without condition for a split, no change asmcfg_bef = asmcfg.copy() -asmcfg.apply_splitting(mdis.symbol_pool) +asmcfg.apply_splitting(mdis.loc_db) assert asmcfg_bef == asmcfg open("graph5.dot", "w").write(asmcfg.dot()) ## Create conditions for a block split -inside_firstbbl = mdis.symbol_pool.getby_offset(4) +inside_firstbbl = mdis.loc_db.get_offset_location(4) tob.bto.add(AsmConstraintTo(inside_firstbbl)) asmcfg.rebuild_edges() assert len(asmcfg.pendings) == 1 assert inside_firstbbl in asmcfg.pendings -asmcfg.apply_splitting(mdis.symbol_pool) +asmcfg.apply_splitting(mdis.loc_db) ## Check result assert len(asmcfg) == 6 assert len(asmcfg.pendings) == 0 @@ -289,7 +289,7 @@ lbl_newb = asmcfg.successors(entry_block.loc_key)[0] newb = asmcfg.loc_key_to_block(lbl_newb) assert len(newb.lines) == 2 assert map(str, newb.lines) == ['XOR ECX, ECX', - 'JNZ loc_3'] + 'JNZ loc_key_3'] preds = asmcfg.predecessors(lbl_newb) assert len(preds) == 2 assert entry_block.loc_key in preds @@ -322,10 +322,10 @@ solutions = list(matcher.match(asmcfg)) assert len(solutions) == 1 solution = solutions.pop() for jbbl, label in solution.iteritems(): - offset = mdis.symbol_pool.loc_key_to_offset(label) + offset = mdis.loc_db.get_location_offset(label) assert offset == int(jbbl._name, 16) -loc_key_dum = mdis.symbol_pool.getby_name_create("dummy_loc") +loc_key_dum = mdis.loc_db.get_or_create_name_location("dummy_loc") asmcfg.add_node(loc_key_dum) error_raised = False try: diff --git a/test/core/locationdb.py b/test/core/locationdb.py new file mode 100644 index 00000000..b9a5f707 --- /dev/null +++ b/test/core/locationdb.py @@ -0,0 +1,108 @@ +from miasm2.core.locationdb import LocationDB + + +# Basic tests (LocationDB description) +loc_db = LocationDB() +loc_key1 = loc_db.add_location() +loc_key2 = loc_db.add_location(offset=0x1234) +loc_key3 = loc_db.add_location(name="first_name") +loc_db.add_location_name(loc_key3, "second_name") +loc_db.set_location_offset(loc_key3, 0x5678) +loc_db.remove_location_name(loc_key3, "second_name") + +assert loc_db.get_location_offset(loc_key1) is None +assert loc_db.get_location_offset(loc_key2) == 0x1234 + +assert loc_db.pretty_str(loc_key1) == str(loc_key1) +assert loc_db.pretty_str(loc_key2) == "loc_1234" +assert loc_db.pretty_str(loc_key3) == "first_name" +loc_db.consistency_check() + +# Offset manipulation +loc_key4 = loc_db.add_location() +assert loc_db.get_location_offset(loc_key4) is None +loc_db.set_location_offset(loc_key4, 0x1122) +assert loc_db.get_location_offset(loc_key4) == 0x1122 +loc_db.unset_location_offset(loc_key4) +assert loc_db.get_location_offset(loc_key4) is None +try: + loc_db.set_location_offset(loc_key4, 0x1234) + has_raised = False +except KeyError: + has_raised = True +assert has_raised +assert loc_db.get_location_offset(loc_key4) is None +loc_db.set_location_offset(loc_key4, 0x1122) +try: + loc_db.set_location_offset(loc_key4, 0x1123) + has_raised = False +except ValueError: + has_raised = True +assert has_raised +assert loc_db.get_location_offset(loc_key4) == 0x1122 +loc_db.set_location_offset(loc_key4, 0x1123, force=True) +assert loc_db.get_location_offset(loc_key4) == 0x1123 +assert 0x1123 in loc_db.offsets +try: + loc_db.add_location(offset=0x1123) + has_raised = False +except ValueError: + has_raised = True +assert loc_db.add_location(offset=0x1123, strict=False) == loc_key4 +assert loc_db.get_offset_location(0x1123) == loc_key4 +assert loc_db.get_or_create_offset_location(0x1123) == loc_key4 +loc_key4_bis = loc_db.get_or_create_offset_location(0x1144) +assert loc_db.get_offset_location(0x1144) == loc_key4_bis +loc_db.consistency_check() + +# Names manipulation +loc_key5 = loc_db.add_location() +name1 = "name1" +name2 = "name2" +name3 = "name3" +assert len(loc_db.get_location_names(loc_key5)) == 0 +loc_db.add_location_name(loc_key5, name1) +loc_db.add_location_name(loc_key5, name2) +assert name1 in loc_db.names +assert name2 in loc_db.names +assert name1 in loc_db.get_location_names(loc_key5) +assert name2 in loc_db.get_location_names(loc_key5) +assert loc_db.get_name_location(name1) == loc_key5 +loc_db.remove_location_name(loc_key5, name1) +assert name1 not in loc_db.names +assert name1 not in loc_db.get_location_names(loc_key5) +try: + loc_db.remove_location_name(loc_key5, name1) + has_raised = False +except KeyError: + has_raised = True +try: + loc_db.add_location_name(loc_key1, name2) + has_raised = False +except KeyError: + has_raised = True +try: + loc_db.add_location(name=name2) + has_raised = False +except ValueError: + has_raised = True +assert loc_db.add_location(name=name2, strict=False) == loc_key5 +assert loc_db.get_or_create_name_location(name2) == loc_key5 +loc_key5_bis = loc_db.get_or_create_name_location(name3) +assert loc_db.get_name_location(name3) == loc_key5_bis +loc_db.consistency_check() + +# Merge +loc_db2 = LocationDB() +loc_db2.add_location(offset=0x3344) +loc_db2.add_location(name=name2) +loc_db.merge(loc_db2) +assert 0x3344 in loc_db.offsets +assert name2 in loc_db.names +loc_db.consistency_check() +assert loc_db.get_name_location(name2) == loc_key5 + +# Delete +loc_db.remove_location(loc_key5) +assert loc_db.get_name_location(name2) is None +loc_db.consistency_check() diff --git a/test/core/parse_asm.py b/test/core/parse_asm.py index fab3a815..ddb195d2 100755 --- a/test/core/parse_asm.py +++ b/test/core/parse_asm.py @@ -64,15 +64,15 @@ class TestParseAsm(unittest.TestCase): .string "toto" ''' - asmcfg, symbol_pool = parse_txt(mn_x86, 32, ASM0) + asmcfg, loc_db = parse_txt(mn_x86, 32, ASM0) patches = asm_resolve_final(mn_x86, asmcfg, - symbol_pool) + loc_db) lbls = [] for i in xrange(6): - lbls.append(symbol_pool.getby_name('lbl%d' % i)) + lbls.append(loc_db.get_name_location('lbl%d' % i)) # align test - offset = symbol_pool.loc_key_to_offset(lbls[5]) + offset = loc_db.get_location_offset(lbls[5]) assert(offset % 0x10 == 0) lbl2block = {} for block in asmcfg.blocks: @@ -95,10 +95,10 @@ class TestParseAsm(unittest.TestCase): RET ''' - asmcfg, symbol_pool = parse_txt(mn_x86, 32, ASM0) + asmcfg, loc_db = parse_txt(mn_x86, 32, ASM0) lbls = [] for i in xrange(2): - lbls.append(symbol_pool.getby_name('lbl%d' % i)) + lbls.append(loc_db.get_name_location('lbl%d' % i)) lbl2block = {} for block in asmcfg.blocks: lbl2block[block.loc_key] = block diff --git a/test/core/sembuilder.py b/test/core/sembuilder.py index 53aa199d..f7a96b89 100644 --- a/test/core/sembuilder.py +++ b/test/core/sembuilder.py @@ -2,15 +2,15 @@ import inspect from pdb import pm from miasm2.core.sembuilder import SemBuilder -from miasm2.core.asmblock import AsmSymbolPool +from miasm2.core.locationdb import LocationDB import miasm2.expression.expression as m2_expr # Test classes class IR(object): - def __init__(self, symbol_pool): - self.symbol_pool = symbol_pool + def __init__(self, loc_db): + self.loc_db = loc_db IRDst = m2_expr.ExprId("IRDst", 32) @@ -45,8 +45,8 @@ def test(Arg1, Arg2, Arg3): a = m2_expr.ExprId('A', 32) b = m2_expr.ExprId('B', 32) c = m2_expr.ExprId('C', 32) -symbol_pool = AsmSymbolPool() -ir = IR(symbol_pool) +loc_db = LocationDB() +ir = IR(loc_db) instr = Instr() res = test(ir, instr, a, b, c) |