1import os
2import sys
3import logging
4
5import angr
6import nose
7
8from common import bin_location
9from test_tracer import tracer_cgc
10
11def test_cgc():
12    binary = os.path.join(bin_location, 'tests', 'cgc', 'sc1_0b32aa01_01')
13    simgr, tracer = tracer_cgc(binary, 'driller_core_cgc', b'AAAA', copy_states=True)
14    simgr.use_technique(angr.exploration_techniques.DrillerCore(tracer._trace))
15    simgr.run()
16
17    nose.tools.assert_true('diverted' in simgr.stashes)
18    nose.tools.assert_equal(len(simgr.diverted), 3)
19
20def test_simprocs():
21    binary = os.path.join(bin_location, 'tests', 'i386', 'driller_simproc')
22    memcmp = angr.SIM_PROCEDURES['libc']['memcmp']()
23
24    simgr, tracer = tracer_cgc(binary, 'driller_core_simprocs', b'A'*128, copy_states=True)
25    p = simgr._project
26    p.hook(0x8048200, memcmp)
27
28    d = angr.exploration_techniques.DrillerCore(tracer._trace)
29    simgr.use_technique(d)
30
31    simgr.run()
32    nose.tools.assert_in('diverted', simgr.stashes)
33    nose.tools.assert_greater(len(simgr.diverted), 0)
34
35
36def run_all():
37    functions = globals()
38    all_functions = dict(filter((lambda kv: kv[0].startswith('test_')), functions.items()))
39    for f in sorted(all_functions.keys()):
40        if hasattr(all_functions[f], '__call__'):
41            all_functions[f]()
42
43
44if __name__ == "__main__":
45    logging.getLogger("angr.exploration_techniques.driller_core").setLevel('DEBUG')
46    if len(sys.argv) > 1:
47        globals()['test_' + sys.argv[1]]()
48    else:
49        run_all()
50