1import os
2import struct
3
4import nose
5import angr
6import angr.simos.windows
7
8
9test_location = os.path.join(os.path.dirname(os.path.realpath(str(__file__))), '..', '..', 'binaries', 'tests')
10
11
12def compare_none(state, test_value):
13    test_value = test_value.concrete
14    correct_value = angr.simos.windows.VS_SECURITY_COOKIES[state.arch.name].default
15    nose.tools.assert_equal(test_value, correct_value)
16
17
18def compare_random(state, test_value):
19    test_value = test_value.concrete
20    incorrect_value = angr.simos.windows.VS_SECURITY_COOKIES[state.arch.name].default
21    nose.tools.assert_not_equal(test_value, incorrect_value)
22
23
24def compare_static(state, test_value):
25    test_value = test_value.concrete
26    correct_value = struct.unpack('>I', b'cook')[0]
27    nose.tools.assert_equal(test_value, correct_value)
28
29
30def compare_symbolic(state, test_value):
31    nose.tools.assert_true(test_value.resolved.symbolic)
32
33
34def check_value(project, init_type, comparison):
35    main_object = project.loader.main_object
36    state = project.factory.blank_state(security_cookie_init=init_type)
37    value = getattr(state.mem[main_object.load_config['SecurityCookie']], "uint{0}_t".format(state.arch.bits))
38    comparison(state, value)
39
40
41def test_security_cookie_init():
42    project = angr.Project(os.path.join(test_location, 'i386', 'test_arrays.exe'))
43    check_value(project, angr.simos.windows.SecurityCookieInit.NONE, compare_none)
44    check_value(project, angr.simos.windows.SecurityCookieInit.RANDOM, compare_random)
45    check_value(project, angr.simos.windows.SecurityCookieInit.STATIC, compare_static)
46    check_value(project, angr.simos.windows.SecurityCookieInit.SYMBOLIC, compare_symbolic)
47
48    nose.tools.assert_raises(TypeError, project.factory.blank_state, security_cookie_init=1)
49
50
51if __name__ == '__main__':
52    test_security_cookie_init()
53