diff --git a/angrmanagement/data/instance.py b/angrmanagement/data/instance.py index 73fbee30bc..c58523493f 100644 --- a/angrmanagement/data/instance.py +++ b/angrmanagement/data/instance.py @@ -14,7 +14,6 @@ from angrmanagement.data.trace import Trace from angrmanagement.errors import ContainerAlreadyRegisteredError from angrmanagement.logic.debugger import DebuggerListManager, DebuggerManager -from angrmanagement.logic.jobmanager import JobManager from .log import LogRecord, initialize from .object_container import ObjectContainer @@ -38,8 +37,6 @@ class Instance: cfb: angr.analyses.cfg.CFBlanket | ObjectContainer log: list[LogRecord] | ObjectContainer - job_manager: JobManager - def __init__(self) -> None: # pylint:disable=import-outside-toplevel # delayed import @@ -50,8 +47,6 @@ def __init__(self) -> None: SavedInteraction, ) - self.job_manager = JobManager(self) - self._live = False self.variable_recovery_job: VariableRecoveryJob | None = None self._analysis_configuration = None @@ -287,7 +282,7 @@ def set_comment(self, addr: int, comment_text) -> None: # TODO: can this be removed? if self.set_comment_callback is not None: - self.set_comment_callback(addr, comment_text) + self.set_comment_callback(addr, comment_text) # pylint:disable=not-callable # # Private methods diff --git a/angrmanagement/logic/debugger/simgr.py b/angrmanagement/logic/debugger/simgr.py index 6e85a91132..05f2021d92 100644 --- a/angrmanagement/logic/debugger/simgr.py +++ b/angrmanagement/logic/debugger/simgr.py @@ -84,7 +84,7 @@ def continue_forward(self) -> None: @property def _num_active_explore_jobs(self) -> int: - return functools.reduce(lambda s, j: s + isinstance(j, SimgrExploreJob), self.instance.job_manager.jobs, 0) + return functools.reduce(lambda s, j: s + isinstance(j, SimgrExploreJob), self.workspace.job_manager.jobs, 0) @property def is_halted(self) -> bool: @@ -95,6 +95,6 @@ def can_halt(self) -> bool: return not self.is_halted def halt(self) -> None: - for job in self.instance.job_manager.jobs: + for job in self.workspace.job_manager.jobs: if isinstance(job, SimgrExploreJob): - self.workspace.main_instance.job_manager.cancel_job(job) + self.workspace.job_manager.cancel_job(job) diff --git a/angrmanagement/logic/jobmanager.py b/angrmanagement/logic/jobmanager.py index 47d2471029..ca999c385c 100644 --- a/angrmanagement/logic/jobmanager.py +++ b/angrmanagement/logic/jobmanager.py @@ -15,8 +15,8 @@ if TYPE_CHECKING: from collections.abc import Callable - from angrmanagement.data.instance import Instance from angrmanagement.data.jobs.job import Job + from angrmanagement.ui.workspace import Workspace log = logging.getLogger(__name__) @@ -78,7 +78,7 @@ def run(self) -> None: log.info('Job "%s" started', job.name) job.start_at = time.time() - result = job.run(ctx, self.job_manager.instance) + result = job.run(ctx, self.job_manager.workspace.main_instance) now = time.time() duration = now - job.start_at log.info('Job "%s" completed after %.2f seconds', job.name, duration) @@ -92,7 +92,7 @@ def run(self) -> None: self.job_manager.job_worker_exception_callback(job, e) else: self.job_manager.jobs.remove(job) - gui_thread_schedule_async(job.finish, args=(self.job_manager.instance, result)) + gui_thread_schedule_async(job.finish, args=(self.job_manager.workspace.main_instance, result)) def keyboard_interrupt(self) -> None: """Called from the GUI thread when the user presses Ctrl+C or presses a cancel button""" @@ -108,7 +108,7 @@ def keyboard_interrupt(self) -> None: class JobManager: """JobManager is responsible for managing jobs and running them in a separate thread.""" - instance: Instance + workspace: Workspace jobs: list[Job] jobs_queue: Queue[Job] @@ -120,8 +120,8 @@ class JobManager: _gui_last_updated_at: float _last_text: str | None - def __init__(self, instance: Instance): - self.instance = instance + def __init__(self, workspace: Workspace): + self.workspace = workspace self.jobs = [] self.jobs_queue = Queue() diff --git a/angrmanagement/plugins/precise_diffing/diff_view.py b/angrmanagement/plugins/precise_diffing/diff_view.py index 62c99250e7..858dd402bb 100644 --- a/angrmanagement/plugins/precise_diffing/diff_view.py +++ b/angrmanagement/plugins/precise_diffing/diff_view.py @@ -103,7 +103,7 @@ def decomp(*args, **kwargs) -> None: # pylint:disable=unused-argument blocking=True, regen_clinic=regen_clinic, ) - self.instance.job_manager.add_job(job) + self.workspace.job_manager.add_job(job) if self._function.ran_cca is False: # run calling convention analysis for this function @@ -113,6 +113,6 @@ def decomp(*args, **kwargs) -> None: # pylint:disable=unused-argument options = {} options["workers"] = 0 varrec_job = VariableRecoveryJob(**options, on_finish=decomp, func_addr=self._function.addr) - self.instance.job_manager.add_job(varrec_job) + self.workspace.job_manager.add_job(varrec_job) else: decomp() diff --git a/angrmanagement/plugins/precise_diffing/precisediff_plugin.py b/angrmanagement/plugins/precise_diffing/precisediff_plugin.py index e7e4615fbc..d40e3d8178 100644 --- a/angrmanagement/plugins/precise_diffing/precisediff_plugin.py +++ b/angrmanagement/plugins/precise_diffing/precisediff_plugin.py @@ -276,11 +276,11 @@ def _create_instance_from_binary(self, file_path: Path) -> None: job = LoadBinaryJob(file_path, on_finish=self._create_instance_from_binary_done) self.loaded_binary = file_path - self.diff_instance.job_manager.add_job(job) + self.workspace.job_manager.add_job(job) def _create_instance_from_binary_done(self, *args, **kwargs) -> None: # pylint:disable=unused-argument job = CFGGenerationJob(on_finish=self._generate_binary_cfg_done) - self.diff_instance.job_manager.add_job(job) + self.workspace.job_manager.add_job(job) def _generate_binary_cfg_done(self, inst, cfg_info, *args, **kwargs) -> None: # pylint:disable=unused-argument cfg_model, _ = cfg_info diff --git a/angrmanagement/ui/main_window.py b/angrmanagement/ui/main_window.py index 32b0b87cab..24c9b018dc 100644 --- a/angrmanagement/ui/main_window.py +++ b/angrmanagement/ui/main_window.py @@ -287,9 +287,7 @@ def _init_statusbar(self) -> None: layout.addWidget(self._stopwatch_label) self._interrupt_job_button = QIconLabel(qta.icon("fa5s.times-circle", color=Conf.palette_buttontext)) - self._interrupt_job_button.clicked.connect( - lambda: self.workspace.main_instance.job_manager.interrupt_current_job - ) + self._interrupt_job_button.clicked.connect(lambda: self.workspace.job_manager.interrupt_current_job) self._interrupt_job_button.hide() layout.addWidget(self._interrupt_job_button) @@ -318,9 +316,9 @@ def _init_statusbar(self) -> None: def on_cancel() -> None: if self.workspace is None: return - for job in self.workspace.main_instance.job_manager.jobs: + for job in self.workspace.job_manager.jobs: if job.blocking: - self.workspace.main_instance.job_manager.interrupt_current_job() + self.workspace.job_manager.interrupt_current_job() break self._progress_dialog.canceled.connect(on_cancel) @@ -401,7 +399,7 @@ def _init_shortcuts(self) -> None: QShortcut(QKeySequence(f"Alt+{i}"), self, lambda idx=i: self._raise_view(idx - 1)) QShortcut(QKeySequence("Alt+0"), self, lambda: self._raise_view(9)) - QShortcut(QKeySequence("Ctrl+I"), self, self.workspace.main_instance.job_manager.interrupt_current_job) + QShortcut(QKeySequence("Ctrl+I"), self, self.workspace.job_manager.interrupt_current_job) # Raise the DisassemblyView after everything has initialized self._raise_view(0) @@ -680,7 +678,7 @@ def open_docker_button(self) -> None: if img_name is None: return target = archr.targets.DockerImageTarget(img_name, target_path=None) - self.workspace.main_instance.job_manager.add_job(LoadTargetJob(target)) + self.workspace.job_manager.add_job(LoadTargetJob(target)) self.workspace.main_instance.img_name = img_name def load_trace_file(self, file_path) -> None: @@ -729,7 +727,7 @@ def load_file(self, file_path) -> None: self._load_database(file_path) else: self._recent_file(file_path) - self.workspace.main_instance.job_manager.add_job(LoadBinaryJob(file_path)) + self.workspace.job_manager.add_job(LoadBinaryJob(file_path)) else: QMessageBox.critical( self, @@ -854,7 +852,7 @@ def run_dependency_analysis(self, func_addr: int | None = None, func_arg_idx: in if self.workspace is None or self.workspace.main_instance is None: return dep_analysis_job = DependencyAnalysisJob(func_addr=func_addr, func_arg_idx=func_arg_idx) - self.workspace.main_instance.job_manager.add_job(dep_analysis_job) + self.workspace.job_manager.add_job(dep_analysis_job) def run_analysis(self) -> None: if self.workspace: @@ -956,7 +954,7 @@ def _load_database(self, file_path: str) -> None: job = LoadAngrDBJob(file_path, ["global", "pseudocode_variable_kb"], other_kbs=other_kbs, extra_info=extra_info) job._on_finish = partial(self._on_load_database_finished, job) - self.workspace.main_instance.job_manager.add_job(job) + self.workspace.job_manager.add_job(job) def _on_load_database_finished(self, job: LoadAngrDBJob, *args, **kwargs) -> None: # pylint:disable=unused-argument proj = job.project diff --git a/angrmanagement/ui/views/code_view.py b/angrmanagement/ui/views/code_view.py index 24fc2370c5..f3aa34c655 100644 --- a/angrmanagement/ui/views/code_view.py +++ b/angrmanagement/ui/views/code_view.py @@ -157,7 +157,7 @@ def decomp(*args, **kwargs) -> None: # pylint:disable=unused-argument blocking=True, regen_clinic=regen_clinic, ) - self.instance.job_manager.add_job(job) + self.workspace.job_manager.add_job(job) if self._function.ran_cca is False: # run calling convention analysis for this function @@ -167,7 +167,7 @@ def decomp(*args, **kwargs) -> None: # pylint:disable=unused-argument options = {} options["workers"] = 0 varrec_job = VariableRecoveryJob(**options, on_finish=decomp, func_addr=self._function.addr) - self.instance.job_manager.add_job(varrec_job) + self.workspace.job_manager.add_job(varrec_job) else: decomp() diff --git a/angrmanagement/ui/widgets/qsimulation_managers.py b/angrmanagement/ui/widgets/qsimulation_managers.py index 515d4ca338..df61f891a1 100644 --- a/angrmanagement/ui/widgets/qsimulation_managers.py +++ b/angrmanagement/ui/widgets/qsimulation_managers.py @@ -245,7 +245,7 @@ def _init_finds_tab(self, tab) -> None: def _on_step_clicked(self) -> None: if not self.simgr.am_none: - self.instance.job_manager.add_job( + self.workspace.job_manager.add_job( SimgrStepJob.create( self.simgr.am_obj, until_branch=False, step_callback=self.workspace.plugins.step_callback ) @@ -253,7 +253,7 @@ def _on_step_clicked(self) -> None: def _on_step_until_branch_clicked(self) -> None: if not self.simgr.am_none: - self.instance.job_manager.add_job( + self.workspace.job_manager.add_job( SimgrStepJob.create( self.simgr.am_obj, until_branch=True, step_callback=self.workspace.plugins.step_callback ) @@ -269,7 +269,7 @@ def _step_callback(simgr): gui_thread_schedule(lambda: self.simgr.am_event(src="post_step")) return simgr - self.instance.job_manager.add_job( + self.workspace.job_manager.add_job( SimgrExploreJob.create( self.simgr, avoid=self.avoid_addrs, find=self.find_addrs, step_callback=_step_callback ) diff --git a/angrmanagement/ui/workspace.py b/angrmanagement/ui/workspace.py index 2f9dc62723..919642949f 100644 --- a/angrmanagement/ui/workspace.py +++ b/angrmanagement/ui/workspace.py @@ -39,6 +39,7 @@ from angrmanagement.logic.debugger import DebuggerWatcher from angrmanagement.logic.debugger.bintrace import BintraceDebugger from angrmanagement.logic.debugger.simgr import SimulationDebugger +from angrmanagement.logic.jobmanager import JobManager from angrmanagement.logic.threads import gui_thread_schedule_async from angrmanagement.plugins import PluginManager from angrmanagement.ui.dialogs import AnalysisOptionsDialog @@ -89,11 +90,15 @@ class Workspace: This class implements the angr management workspace. """ + job_manager: JobManager + def __init__(self, main_window: MainWindow, instance: Instance) -> None: self.main_window: MainWindow = main_window self._main_instance = instance instance.workspace = self + self.job_manager = JobManager(self) + self.command_manager: CommandManager = CommandManager() self.view_manager: ViewManager = ViewManager(self) self.plugins: PluginManager = PluginManager(self) @@ -218,7 +223,7 @@ def generate_cfg(self, cfg_args=None) -> None: cfg_args = {} cfg_job = CFGGenerationJob(on_finish=self.on_cfg_generated, **cfg_args) - self.main_instance.job_manager.add_job(cfg_job) + self.job_manager.add_job(cfg_job) start_daemon_thread(self._refresh_cfg, "Progressively Refreshing CFG", args=(cfg_job,)) def _refresh_cfg(self, cfg_job) -> None: @@ -245,7 +250,7 @@ def _refresh_cfg(self, cfg_job) -> None: reloaded = True time.sleep(0.3) - if cfg_job not in self.main_instance.job_manager.jobs: + if cfg_job not in self.job_manager.jobs: break def on_cfg_generated(self, instance, cfg_result) -> None: # pylint:disable=unused-argument @@ -256,7 +261,7 @@ def on_cfg_generated(self, instance, cfg_result) -> None: # pylint:disable=unus self.main_instance.cfg.am_event() if self.main_instance._analysis_configuration["flirt"].enabled: - self.main_instance.job_manager.add_job( + self.job_manager.add_job( FlirtSignatureRecognitionJob( on_finish=self._on_flirt_signature_recognized, ) @@ -287,7 +292,7 @@ def on_cfg_generated(self, instance, cfg_result) -> None: # pylint:disable=unus view.clear() def _on_flirt_signature_recognized(self, *args, **kwargs) -> None: # pylint:disable=unused-argument - self.main_instance.job_manager.add_job( + self.job_manager.add_job( PrototypeFindingJob( on_finish=self._on_prototype_found, ) @@ -295,7 +300,7 @@ def _on_flirt_signature_recognized(self, *args, **kwargs) -> None: # pylint:dis def _on_prototype_found(self, *args, **kwargs) -> None: # pylint:disable=unused-argument if self.main_instance._analysis_configuration["code_tagging"].enabled: - self.main_instance.job_manager.add_job( + self.job_manager.add_job( CodeTaggingJob( on_finish=self.on_function_tagged, ) @@ -314,7 +319,7 @@ def _on_prototype_found(self, *args, **kwargs) -> None: # pylint:disable=unused disassembly_view = self.view_manager.first_view_in_category("disassembly") if disassembly_view is not None and not disassembly_view.function.am_none: self.main_instance.variable_recovery_job.prioritize_function(disassembly_view.function.addr) - self.main_instance.job_manager.add_job(self.main_instance.variable_recovery_job) + self.job_manager.add_job(self.main_instance.variable_recovery_job) def _on_patch_event(self, **kwargs) -> None: if self.main_instance.cfg.am_none: @@ -729,7 +734,7 @@ def create_project_from_trace(self, trace: Trace, on_complete: Callable) -> None self.main_instance.binary_path = thing self.main_instance.original_binary_path = thing job = LoadBinaryJob(thing, load_options=load_options, on_finish=on_complete) - self.main_instance.job_manager.add_job(job) + self.job_manager.add_job(job) def interact_program(self, img_name: str, view=None) -> None: if view is None or view.category != "interaction": diff --git a/tests/common.py b/tests/common.py index 29fd7bde99..9f84d7c9b0 100644 --- a/tests/common.py +++ b/tests/common.py @@ -110,4 +110,4 @@ def setUp(self): os.path.join(test_location, "x86_64", "true"), auto_load_libs=False ) self.main.workspace.main_instance.project.am_event() - self.main.workspace.main_instance.job_manager.join_all_jobs() + self.main.workspace.job_manager.join_all_jobs() diff --git a/tests/manual_human_activities.py b/tests/manual_human_activities.py index fbe49a62ae..628d59cd81 100644 --- a/tests/manual_human_activities.py +++ b/tests/manual_human_activities.py @@ -28,7 +28,7 @@ def _open_a_project(self): binpath = os.path.join(test_location, "x86_64", "fauxware") main.workspace.main_instance.project.am_obj = angr.Project(binpath, auto_load_libs=False) main.workspace.main_instance.project.am_event() - main.workspace.main_instance.job_manager.join_all_jobs() + main.workspace.job_manager.join_all_jobs() self.project = binpath # import ipdb; ipdb.set_trace() self.project_md5 = main.workspace.main_instance.project.loader.main_object.md5.hex() @@ -48,7 +48,7 @@ def test_rename_a_function_in_disasm_and_pseudocode_views(self): disasm_view.display_disasm_graph() disasm_view.display_function(func) disasm_view.decompile_current_function() - main.workspace.main_instance.job_manager.join_all_jobs() + main.workspace.job_manager.join_all_jobs() pseudocode_view = main.workspace._get_or_create_view("pseudocode", CodeView) # find the node for function @@ -89,7 +89,7 @@ def test_rename_a_variable_in_pseudocode_view(self): disasm_view.display_disasm_graph() disasm_view.display_function(func) disasm_view.decompile_current_function() - main.workspace.main_instance.job_manager.join_all_jobs() + main.workspace.job_manager.join_all_jobs() pseudocode_view = main.workspace._get_or_create_view("pseudocode", CodeView) # find an arbitrary node for a variable diff --git a/tests/test_rename_functions.py b/tests/test_rename_functions.py index 987f9e626b..4ac36bab60 100644 --- a/tests/test_rename_functions.py +++ b/tests/test_rename_functions.py @@ -66,7 +66,7 @@ def test_rename_a_function_in_disasm_and_pseudocode_views(self): binpath = os.path.join(test_location, "x86_64", "fauxware") main.workspace.main_instance.project.am_obj = angr.Project(binpath, auto_load_libs=False) main.workspace.main_instance.project.am_event() - main.workspace.main_instance.job_manager.join_all_jobs() + main.workspace.job_manager.join_all_jobs() func = main.workspace.main_instance.project.kb.functions["main"] self.assertIsNotNone(func) @@ -76,7 +76,7 @@ def test_rename_a_function_in_disasm_and_pseudocode_views(self): disasm_view.display_disasm_graph() gui_thread_schedule(disasm_view.display_function, args=(func,)) disasm_view.decompile_current_function() - main.workspace.main_instance.job_manager.join_all_jobs() + main.workspace.job_manager.join_all_jobs() # run the jobless method in the GUI thread gui_thread_schedule(self._test_rename_a_function_in_disasm_and_pseudocode_views) @@ -109,7 +109,7 @@ def test_rename_a_callee_in_pseudocode_view(self): binpath = os.path.join(test_location, "x86_64", "fauxware") main.workspace.main_instance.project.am_obj = angr.Project(binpath, auto_load_libs=False) main.workspace.main_instance.project.am_event() - main.workspace.main_instance.job_manager.join_all_jobs() + main.workspace.job_manager.join_all_jobs() func = main.workspace.main_instance.project.kb.functions["main"] self.assertIsNotNone(func) @@ -119,7 +119,7 @@ def test_rename_a_callee_in_pseudocode_view(self): disasm_view.display_disasm_graph() gui_thread_schedule(disasm_view.display_function, args=(func,)) disasm_view.decompile_current_function() - main.workspace.main_instance.job_manager.join_all_jobs() + main.workspace.job_manager.join_all_jobs() # run the jobless method in the GUI thread gui_thread_schedule(self._test_rename_a_callee_in_pseudocode_view) diff --git a/tests/test_rename_variables.py b/tests/test_rename_variables.py index 02ac235f44..59b03ed07d 100644 --- a/tests/test_rename_variables.py +++ b/tests/test_rename_variables.py @@ -29,7 +29,7 @@ def setUp(self) -> None: proj = angr.Project(binpath, auto_load_libs=False) self.main.workspace.main_instance.project.am_obj = proj self.main.workspace.main_instance.project.am_event() - self.main.workspace.main_instance.job_manager.join_all_jobs() + self.main.workspace.job_manager.join_all_jobs() self.func = proj.kb.functions["doit"] self.assertIsNotNone(self.func) @@ -39,7 +39,7 @@ def setUp(self) -> None: disasm_view.display_disasm_graph() gui_thread_schedule(disasm_view.display_function, args=(self.func,)) disasm_view.decompile_current_function() - self.main.workspace.main_instance.job_manager.join_all_jobs() + self.main.workspace.job_manager.join_all_jobs() self.code_view: CodeView = self.main.workspace.view_manager.first_view_in_category("pseudocode") def tearDown(self) -> None: diff --git a/tests/test_workflow.py b/tests/test_workflow.py index ea85f2593a..3c3c4fe11f 100644 --- a/tests/test_workflow.py +++ b/tests/test_workflow.py @@ -24,7 +24,7 @@ def test_workflow(self): proj = angr.Project(os.path.join(test_location, "x86_64", "true"), auto_load_libs=False) main.workspace.main_instance.project.am_obj = proj main.workspace.main_instance.project.am_event() - main.workspace.main_instance.job_manager.join_all_jobs() + main.workspace.job_manager.join_all_jobs() if __name__ == "__main__":