1import nose
2import os
3import unittest
4
5from archinfo import ArchAMD64
6
7import angr
8from angr.utils.constants import DEFAULT_STATEMENT
9
10TEST_LOCATION = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', '..', 'binaries', 'tests')
11
12
13class TestFunctionManager(unittest.TestCase):
14    @classmethod
15    def setUpClass(cls):
16        cls.project = angr.Project(os.path.join(TEST_LOCATION, "x86_64", "fauxware"))
17
18
19    def test_amd64(self):
20        expected_functions = { 0x4004e0, 0x400510, 0x400520, 0x400530, 0x400540, 0x400550, 0x400560, 0x400570,
21                               0x400580, 0x4005ac, 0x400640, 0x400664, 0x4006ed, 0x4006fd, 0x40071d, 0x4007e0,
22                               0x400880 }
23        expected_blocks = { 0x40071D, 0x40073E, 0x400754, 0x40076A, 0x400774, 0x40078A, 0x4007A0, 0x4007B3, 0x4007C7,
24                            0x4007C9, 0x4007BD, 0x4007D3 }
25        expected_callsites = { 0x40071D, 0x40073E, 0x400754, 0x40076A, 0x400774, 0x40078A, 0x4007A0, 0x4007BD, 0x4007C9 }
26        expected_callsite_targets = { 4195600, 4195632, 4195632, 4195600, 4195632, 4195632, 4195940, 4196077, 4196093 }
27        expected_callsite_returns = { 0x40073e, 0x400754, 0x40076a, 0x400774, 0x40078a, 0x4007a0, 0x4007b3, 0x4007c7,
28                                      None }
29
30        cfg = self.project.analyses.CFGEmulated()  # pylint:disable=unused-variable
31        nose.tools.assert_equal(
32            { k for k in self.project.kb.functions.keys() if k < 0x500000 },
33            expected_functions
34        )
35
36        main = self.project.kb.functions.function(name='main')
37        nose.tools.assert_equal(main.startpoint.addr, 0x40071D)
38        nose.tools.assert_equal(set(main.block_addrs), expected_blocks)
39        nose.tools.assert_equal([0x4007D3], [bl.addr for bl in main.endpoints])
40        nose.tools.assert_equal(set(main.get_call_sites()), expected_callsites)
41        nose.tools.assert_equal(
42            set(map(main.get_call_target, main.get_call_sites())),
43            expected_callsite_targets
44        )
45        nose.tools.assert_equal(
46            set(map(main.get_call_return, main.get_call_sites())),
47            expected_callsite_returns
48        )
49        nose.tools.assert_true(main.has_return)
50
51        rejected = self.project.kb.functions.function(name='rejected')
52        nose.tools.assert_equal(rejected.returning, False)
53
54        # transition graph
55        main_g = main.transition_graph
56        main_g_edges_ = main_g.edges(data=True)
57
58        # Convert nodes those edges from blocks to addresses
59        main_g_edges = []
60        for src_node, dst_node, data in main_g_edges_:
61            main_g_edges.append((src_node.addr, dst_node.addr, data))
62
63        edges = [
64            (0x40071d, 0x400510, {'type': 'call', 'stmt_idx': DEFAULT_STATEMENT, 'ins_addr': 0x400739}),
65            (0x40071d, 0x400510, {'type': 'call', 'stmt_idx': DEFAULT_STATEMENT, 'ins_addr': 0x400739}),
66            (0x40071d, 0x40073e, {'type': 'fake_return', 'confirmed': True, 'outside': False}),
67            (0x40073e, 0x400530, {'type': 'call', 'stmt_idx': DEFAULT_STATEMENT, 'ins_addr': 0x40074f}),
68            (0x40073e, 0x400754, {'type': 'fake_return', 'confirmed': True, 'outside': False}),
69            # rejected() does not return
70            (0x4007c9, 0x4006fd, {'type': 'call', 'stmt_idx': DEFAULT_STATEMENT, 'ins_addr': 0x4007ce}),
71            (0x4007c9, 0x4007d3, {'type': 'fake_return', 'outside': False}),
72        ]
73        for edge in edges:
74            nose.tools.assert_true(edge in main_g_edges)
75
76        # These tests fail for reasons of fastpath, probably
77        #nose.tools.assert_true(main.bp_on_stack)
78        #nose.tools.assert_equal(main.name, 'main')
79        #nose.tools.assert_true(main.retaddr_on_stack)
80        #nose.tools.assert_equal(0x50, main.sp_difference)
81
82        # TODO: Check the result returned
83        #func_man.dbg_draw()
84
85    def test_call_to(self):
86        self.project.arch = ArchAMD64()
87
88        self.project.kb.functions._add_call_to(0x400000, 0x400410, 0x400420, 0x400414)
89        nose.tools.assert_in(0x400000, self.project.kb.functions.keys())
90        nose.tools.assert_in(0x400420, self.project.kb.functions.keys())
91