diff --git a/bin/deepstate/main_angr.py b/bin/deepstate/main_angr.py index c2020c1..30bf579 100644 --- a/bin/deepstate/main_angr.py +++ b/bin/deepstate/main_angr.py @@ -292,6 +292,21 @@ def run_test(project, test, apis, run_state): L.error("Uncaught exception: {}\n{}".format(e, traceback.format_exc())) +def find_symbol_ea(project, name): + try: + ea = project.kb.labels.lookup(name) + if ea: + return ea + except: + pass + + try: + return project.kb.labels.lookup("_{}".format(name)) + except: + pass + + return 0 + def main(): """Run DeepState.""" args = DeepAngr.parse_args() @@ -314,19 +329,11 @@ def main(): args.binary, e)) return 1 - try: - setup_ea = project.kb.labels.lookup('DeepState_Setup') - if not setup_ea: - raise Exception() - except: - try: - setup_ea = project.kb.labels.lookup('_DeepState_Setup') - if not setup_ea: - raise Exception() - except: - L.critical("Cannot find symbol `DeepState_Setup` in binary `{}`".format( + setup_ea = find_symbol_ea(project, 'DeepState_Setup') + if not setup_ea: + L.critical("Cannot find symbol `DeepState_Setup` in binary `{}`".format( args.binary)) - return 1 + return 1 entry_state = project.factory.entry_state( add_options={angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY, @@ -351,10 +358,10 @@ def main(): # symbols. Technically we can look these up with the `labels.lookup` API, # but we have the API table for Manticore-compatibility, so we may as well # use it. - try: - ea_of_api_table = project.kb.labels.lookup('DeepState_API') - except: - ea_of_api_table = project.kb.labels.lookup('_DeepState_API') + ea_of_api_table = find_symbol_ea(project, 'DeepState_API') + if not ea_of_api_table: + L.critical("Could not find API table in binary `{}`".format(args.binary)) + return 1 mc = DeepAngr(state=run_state) apis = mc.read_api_table(ea_of_api_table) diff --git a/bin/deepstate/main_manticore.py b/bin/deepstate/main_manticore.py index 5b57450..d53eadf 100644 --- a/bin/deepstate/main_manticore.py +++ b/bin/deepstate/main_manticore.py @@ -246,6 +246,22 @@ def done_test(_, state, state_id, reason): mc.report() +def find_symbol_ea(m, name): + try: + ea = m._get_symbol_address(name) + if ea: + return ea + except: + pass + + try: + return m._get_symbol_address("_{}".format(name)) + except: + pass + + return 0 + + def do_run_test(state, apis, test): """Run an individual test case.""" state.cpu.PC = test.ea @@ -322,27 +338,20 @@ def main(): m._binary_type = 'not elf' m._binary_obj = m._initial_state.platform.elf - try: - setup_ea = m._get_symbol_address('DeepState_Setup') - if not setup_ea: - raise Exception("Could not find symbol") - except: - try: - setup_ea = m._get_symbol_address('_DeepState_Setup') - if not setup_ea: - raise Exception("Could not find symbol") - except: - L.critical("Cannot find symbol `DeepState_Setup` in binary `{}`: {}".format( - args.binary, traceback.format_exc())) - return 1 + setup_ea = find_symbol_ea(m, 'DeepState_Setup') + if not setup_ea: + L.critical("Cannot find symbol `DeepState_Setup` in binary `{}`".format( + args.binary)) + return 1 setup_state = m._initial_state mc = DeepManticore(setup_state) - ea_of_api_table = m._get_symbol_address('DeepState_API') + ea_of_api_table = find_symbol_ea(m, 'DeepState_API') if not ea_of_api_table: - ea_of_api_table = m._get_symbol_address('_DeepState_API') + L.critical("Could not find API table in binary `{}`".format(args.binary)) + return 1 apis = mc.read_api_table(ea_of_api_table) del mc