1import os
2import angr
3import nose
4
5
6test_location = str(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../binaries/tests'))
7arches = {'x86_64'}
8
9def main():
10    test_ddg_global_var_dependencies()
11
12def test_ddg_global_var_dependencies():
13    for arch in arches:
14        run_ddg_global_var_dependencies(arch)
15
16def run_ddg_global_var_dependencies(arch):
17    test_file = os.path.join(test_location, arch, 'ddg_global_var_dependencies')
18    proj = angr.Project(test_file, auto_load_libs=False)
19    cfg = proj.analyses.CFGEmulated(context_sensitivity_level=2, keep_state=True, state_add_options=angr.sim_options.refs)
20    ddg = proj.analyses.DDG(cfg)
21    main_func = cfg.functions.function(name='main')
22
23    target_block_addr = main_func.ret_sites[0].addr
24    target_block = proj.factory.block(addr=target_block_addr)
25    tgt_stmt_idx, tgt_stmt = get_target_stmt(proj, target_block)
26    assert tgt_stmt_idx is not None
27    buf_addr = tgt_stmt.data.addr.con.value
28    tgt_ddg_node = get_ddg_node(ddg, target_block_addr, tgt_stmt_idx)
29    assert tgt_ddg_node is not None
30
31    # Whether the target depends on the statement assigning 'b' to the global variable
32    has_correct_dependency = False
33    for pred in ddg.get_predecessors(tgt_ddg_node):
34        pred_block = proj.factory.block(addr=pred.block_addr)
35        stmt = pred_block.vex.statements[pred.stmt_idx]
36        has_correct_dependency |= check_dependency(stmt, buf_addr, ord('b'))
37
38        # If the target depends on the statement assigning 'a' to the global variable, it is underconstrained (this assignment should be overwritten by the 'b' assignment)
39        nose.tools.assert_false(check_dependency(stmt, buf_addr, ord('a')), msg="Target statement has incorrect dependency (DDG is underconstrained)")
40    nose.tools.assert_true(has_correct_dependency, msg='Target statement does not have correct dependency (DDG is overconstrained)')
41
42
43
44def check_dependency(stmt, addr, const):
45    # Check if we are storing a constant to a variable with constant address
46    if stmt.tag == 'Ist_Store' and stmt.addr.tag == 'Iex_Const' and stmt.data.tag == 'Iex_Const':
47        # Check if we are storing the specified constant to the specified variable address
48        if stmt.addr.con.value == addr and stmt.data.con.value == const:
49            return True
50
51    return False
52
53def get_ddg_node(ddg, block_addr, stmt_idx):
54    for node in ddg.graph.nodes:
55        if node.block_addr == block_addr and node.stmt_idx == stmt_idx:
56            return node
57    return None
58
59def get_target_stmt(proj, block):
60    for i, stmt in enumerate(block.vex.statements):
61        # We're looking for the instruction that loads a constant memory address into a temporary variable
62        if stmt.tag == 'Ist_WrTmp' and stmt.data.tag == 'Iex_Load' and stmt.data.addr.tag == 'Iex_Const':
63            addr = stmt.data.addr.con.value
64            section = proj.loader.main_object.find_section_containing(addr)
65            # Confirm the memory address is in the uninitialized data section
66            if section.name == '.bss':
67                return i, stmt
68    return None, None
69
70
71if __name__ == "__main__":
72    main()
73