diff options
Diffstat (limited to 'test')
| -rw-r--r-- | test/analysis/depgraph.py | 80 | ||||
| -rw-r--r-- | test/analysis/dse.py | 2 | ||||
| -rw-r--r-- | test/arch/aarch64/unit/asm_test.py | 2 | ||||
| -rw-r--r-- | test/arch/mips32/unit/asm_test.py | 2 | ||||
| -rwxr-xr-x | test/arch/x86/sem.py | 4 | ||||
| -rw-r--r-- | test/arch/x86/unit/asm_test.py | 2 | ||||
| -rwxr-xr-x | test/arch/x86/unit/mn_strings.py | 12 | ||||
| -rw-r--r-- | test/core/asmblock.py | 30 | ||||
| -rwxr-xr-x | test/core/parse_asm.py | 6 |
9 files changed, 33 insertions, 107 deletions
diff --git a/test/analysis/depgraph.py b/test/analysis/depgraph.py index a7d3a148..f0772cca 100644 --- a/test/analysis/depgraph.py +++ b/test/analysis/depgraph.py @@ -100,79 +100,6 @@ class IRATest(ira): return set([self.ret_reg, self.sp]) -def bloc2graph(irgraph, label=False, lines=True): - """Render dot graph of @blocks""" - - escape_chars = re.compile('[' + re.escape('{}') + ']') - label_attr = 'colspan="2" align="center" bgcolor="grey"' - edge_attr = 'label = "%s" color="%s" style="bold"' - td_attr = 'align="left"' - block_attr = 'shape="Mrecord" fontname="Courier New"' - - out = ["digraph asm_graph {"] - fix_chars = lambda x: '\\' + x.group() - - # Generate basic blocks - out_blocks = [] - for label in irgraph.graph.nodes(): - if isinstance(label, LocKey): - label_name = irgraph.loc_db.loc_key_to_name(label) - else: - label_name = str(label) - - if hasattr(irgraph, 'blocks'): - irblock = irgraph.blocks[label] - else: - irblock = None - if isinstance(label, LocKey): - out_block = '%s [\n' % label_name - else: - out_block = '%s [\n' % label - out_block += "%s " % block_attr - out_block += 'label =<<table border="0" cellborder="0" cellpadding="3">' - - block_label = '<tr><td %s>%s</td></tr>' % ( - label_attr, label_name) - block_html_lines = [] - if lines and irblock is not None: - for assignblk in irblock: - for dst, src in assignblk.iteritems(): - if False: - out_render = "%.8X</td><td %s> " % (0, td_attr) - else: - out_render = "" - out_render += escape_chars.sub(fix_chars, "%s = %s" % (dst, src)) - block_html_lines.append(out_render) - block_html_lines.append(" ") - block_html_lines.pop() - block_html_lines = ('<tr><td %s>' % td_attr + - ('</td></tr><tr><td %s>' % td_attr).join(block_html_lines) + - '</td></tr>') - out_block += "%s " % block_label - out_block += block_html_lines + "</table>> ];" - out_blocks.append(out_block) - - out += out_blocks - # Generate links - for src, dst in irgraph.graph.edges(): - if isinstance(src, LocKey): - src_name = irgraph.loc_db.loc_key_to_name(src) - else: - src_name = str(src) - if isinstance(dst, LocKey): - dst_name = irgraph.loc_db.loc_key_to_name(dst) - else: - dst_name = str(dst) - - edge_color = "black" - out.append('%s -> %s' % (src_name, - dst_name) + - '[' + edge_attr % ("", edge_color) + '];') - - out.append("}") - return '\n'.join(out) - - def dg2graph(graph, label=False, lines=True): """Render dot graph of @blocks""" @@ -189,7 +116,7 @@ def dg2graph(graph, label=False, lines=True): out_blocks = [] for node in graph.nodes(): if isinstance(node, DependencyNode): - name = loc_db.loc_key_to_name(node.loc_key) + name = loc_db.pretty_str(node.loc_key) node_name = "%s %s %s" % (name, node.element, node.line_nb) @@ -642,7 +569,7 @@ def flatNode(node): element = int(node.element.arg) else: RuntimeError("Unsupported type '%s'" % type(enode.element)) - name = loc_db.loc_key_to_name(node.loc_key) + name = loc_db.pretty_str(node.loc_key) return (name, element, node.line_nb) @@ -741,7 +668,7 @@ def match_results(resultsA, resultsB, nodes): def get_flat_init_depnodes(depnodes): out = [] for node in depnodes: - name = loc_db.loc_key_to_name(node.loc_key) + name = loc_db.pretty_str(node.loc_key) out.append((name, node.element.name, node.line_nb, @@ -1026,7 +953,6 @@ for test_nb, test in enumerate([(G1_IRA, G1_INPUT), g_ira, (depnodes, heads) = test open("graph_%02d.dot" % (test_nb + 1), "w").write(g_ira.graph.dot()) - open("graph_%02d.dot" % (test_nb + 1), "w").write(bloc2graph(g_ira)) # Different options suffix_key_list = ["", "_nosimp", "_nomem", "_nocall", diff --git a/test/analysis/dse.py b/test/analysis/dse.py index 56de4d4e..a05d8595 100644 --- a/test/analysis/dse.py +++ b/test/analysis/dse.py @@ -77,7 +77,7 @@ class DSETest(object): ) # fix shellcode addr - loc_db.set_offset(loc_db.getby_name("main"), 0x0) + loc_db.set_location_offset(loc_db.get_name_location("main"), 0x0) output = StrPatchwork() patches = asm_resolve_final(mn_x86, blocks, loc_db) for offset, raw in patches.items(): diff --git a/test/arch/aarch64/unit/asm_test.py b/test/arch/aarch64/unit/asm_test.py index d3221ec0..677d474f 100644 --- a/test/arch/aarch64/unit/asm_test.py +++ b/test/arch/aarch64/unit/asm_test.py @@ -25,7 +25,7 @@ class Asm_Test(object): blocks, loc_db = parse_asm.parse_txt(mn_aarch64, 'l', self.TXT, loc_db = self.myjit.ir_arch.loc_db) # fix shellcode addr - loc_db.set_offset(loc_db.getby_name("main"), 0x0) + loc_db.set_location_offset(loc_db.get_name_location("main"), 0x0) s = StrPatchwork() patches = asmblock.asm_resolve_final(mn_aarch64, blocks, loc_db) for offset, raw in patches.items(): diff --git a/test/arch/mips32/unit/asm_test.py b/test/arch/mips32/unit/asm_test.py index be26bf49..da792874 100644 --- a/test/arch/mips32/unit/asm_test.py +++ b/test/arch/mips32/unit/asm_test.py @@ -27,7 +27,7 @@ class Asm_Test(object): blocks, loc_db = parse_asm.parse_txt(mn_mips32, 'l', self.TXT, loc_db=self.myjit.ir_arch.loc_db) # fix shellcode addr - loc_db.set_offset(loc_db.getby_name("main"), 0x0) + loc_db.set_location_offset(loc_db.get_name_location("main"), 0x0) s = StrPatchwork() patches = asmblock.asm_resolve_final(mn_mips32, blocks, loc_db) for offset, raw in patches.items(): diff --git a/test/arch/x86/sem.py b/test/arch/x86/sem.py index 23b22245..a2493d12 100755 --- a/test/arch/x86/sem.py +++ b/test/arch/x86/sem.py @@ -49,10 +49,10 @@ def compute(ir, mode, asm, inputstate={}, debug=False): def compute_txt(ir, mode, txt, inputstate={}, debug=False): asmcfg, loc_db = parse_asm.parse_txt(mn, mode, txt) - loc_db.set_offset(loc_db.getby_name("main"), 0x0) + loc_db.set_location_offset(loc_db.get_name_location("main"), 0x0) patches = asmblock.asm_resolve_final(mn, asmcfg, loc_db) interm = ir(loc_db) - lbl = loc_db.getby_name("main") + lbl = loc_db.get_name_location("main") for bbl in asmcfg.blocks: interm.add_block(bbl) return symb_exec(lbl, interm, inputstate, debug) diff --git a/test/arch/x86/unit/asm_test.py b/test/arch/x86/unit/asm_test.py index e626768d..91da1942 100644 --- a/test/arch/x86/unit/asm_test.py +++ b/test/arch/x86/unit/asm_test.py @@ -43,7 +43,7 @@ class Asm_Test(object): blocks, loc_db = parse_asm.parse_txt(mn_x86, self.arch_attrib, self.TXT, loc_db = self.myjit.ir_arch.loc_db) # fix shellcode addr - loc_db.set_offset(loc_db.getby_name("main"), 0x0) + loc_db.set_location_offset(loc_db.get_name_location("main"), 0x0) s = StrPatchwork() patches = asmblock.asm_resolve_final(mn_x86, blocks, loc_db) for offset, raw in patches.items(): diff --git a/test/arch/x86/unit/mn_strings.py b/test/arch/x86/unit/mn_strings.py index a6facd08..8ca148e5 100755 --- a/test/arch/x86/unit/mn_strings.py +++ b/test/arch/x86/unit/mn_strings.py @@ -21,8 +21,8 @@ class Test_SCAS(Asm_Test_32): def check(self): assert(self.myjit.cpu.ECX == len(self.MYSTRING)) - mystr = self.myjit.ir_arch.loc_db.getby_name('mystr') - assert(self.myjit.cpu.EDI == self.myjit.ir_arch.loc_db.loc_key_to_offset(mystr) + len(self.MYSTRING)+1) + mystr = self.myjit.ir_arch.loc_db.get_name_location('mystr') + assert(self.myjit.cpu.EDI == self.myjit.ir_arch.loc_db.get_location_offset(mystr) + len(self.MYSTRING)+1) class Test_MOVS(Asm_Test_32): @@ -43,10 +43,10 @@ class Test_MOVS(Asm_Test_32): def check(self): assert(self.myjit.cpu.ECX == 0) - buffer = self.myjit.ir_arch.loc_db.getby_name('buffer') - assert(self.myjit.cpu.EDI == self.myjit.ir_arch.loc_db.loc_key_to_offset(buffer) + len(self.MYSTRING)) - mystr = self.myjit.ir_arch.loc_db.getby_name('mystr') - assert(self.myjit.cpu.ESI == self.myjit.ir_arch.loc_db.loc_key_to_offset(mystr) + len(self.MYSTRING)) + buffer = self.myjit.ir_arch.loc_db.get_name_location('buffer') + assert(self.myjit.cpu.EDI == self.myjit.ir_arch.loc_db.get_location_offset(buffer) + len(self.MYSTRING)) + mystr = self.myjit.ir_arch.loc_db.get_name_location('mystr') + assert(self.myjit.cpu.ESI == self.myjit.ir_arch.loc_db.get_location_offset(mystr) + len(self.MYSTRING)) if __name__ == "__main__": diff --git a/test/core/asmblock.py b/test/core/asmblock.py index 465697f4..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.loc_db.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.loc_db.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.loc_db.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.loc_db.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.loc_db.getby_name_create("testlabel_src1") -testlabel_src2 = mdis.loc_db.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.loc_db.getby_name_create("testlabel_pend_src") -testlabel_pend_dst = mdis.loc_db.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,11 +258,11 @@ 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 @@ -272,7 +272,7 @@ 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.loc_db.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 @@ -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.loc_db.loc_key_to_offset(label) + offset = mdis.loc_db.get_location_offset(label) assert offset == int(jbbl._name, 16) -loc_key_dum = mdis.loc_db.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/parse_asm.py b/test/core/parse_asm.py index e8d5adc5..ddb195d2 100755 --- a/test/core/parse_asm.py +++ b/test/core/parse_asm.py @@ -70,9 +70,9 @@ class TestParseAsm(unittest.TestCase): loc_db) lbls = [] for i in xrange(6): - lbls.append(loc_db.getby_name('lbl%d' % i)) + lbls.append(loc_db.get_name_location('lbl%d' % i)) # align test - offset = loc_db.loc_key_to_offset(lbls[5]) + offset = loc_db.get_location_offset(lbls[5]) assert(offset % 0x10 == 0) lbl2block = {} for block in asmcfg.blocks: @@ -98,7 +98,7 @@ class TestParseAsm(unittest.TestCase): asmcfg, loc_db = parse_txt(mn_x86, 32, ASM0) lbls = [] for i in xrange(2): - lbls.append(loc_db.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 |