Scan non-blockingly

Author: juliangrtz

core/concurrency/subprocess.py

@@ -0,0 +1,44 @@
+from concurrent.futures import ProcessPoolExecutor, Future
+from functools import partial
+from typing import Callable, Optional, Any
+
+
+# Used for heavy calculations where a thread is not enough and blocking the UI is not feasible.
+class Subprocess:
+    def __init__(self):
+        self.executor = ProcessPoolExecutor()
+        self.future: Optional[Future] = None
+
+    def submit(
+            self,
+            func: Callable,
+            *args,
+            on_done: Optional[Callable[[Any], None]] = None,
+            on_error: Optional[Callable[[Exception], None]] = None,
+            **kwargs
+    ):
+        if self.future and not self.future.done():
+            self.future.cancel()
+
+        self.future = self.executor.submit(func, *args, **kwargs)
+        self.future.add_done_callback(
+            partial(self._handle_future, on_done=on_done, on_error=on_error)
+        )
+
+    def _handle_future(self, future: Future, on_done: Optional[Callable], on_error: Optional[Callable]):
+        self._process_future(future, on_done, on_error)
+
+    def _process_future(self, future: Future, on_done: Optional[Callable], on_error: Optional[Callable]):
+        try:
+            result = future.result()
+        except Exception as e:
+            if on_error:
+                on_error(e)
+        else:
+            if on_done:
+                on_done(result)
+
+    def shutdown(self, wait: bool = False, cancel_futures: bool = True):
+        if self.future and not self.future.done():
+            self.future.cancel()
+        self.executor.shutdown(wait=wait, cancel_futures=cancel_futures)

core/concurrency/worker.py

@@ -0,0 +1,34 @@
+import sys
+import traceback
+
+from PySide6.QtCore import QObject, Signal, QRunnable, Slot
+
+
+# Used for lighter calculations where a thread is enough.
+# Briefly blocking the UI is acceptable.
+class WorkerSignals(QObject):
+    finished = Signal()
+    error = Signal(tuple)
+    result = Signal(object)
+
+
+class Worker(QRunnable):
+    def __init__(self, fn, *args, **kwargs):
+        super().__init__()
+        self.fn = fn
+        self.args = args
+        self.kwargs = kwargs
+        self.signals = WorkerSignals()
+
+    @Slot()
+    def run(self):
+        try:
+            result = self.fn(*self.args, **self.kwargs)
+        except Exception:
+            traceback.print_exc()
+            exctype, value = sys.exc_info()[:2]
+            self.signals.error.emit((exctype, value, traceback.format_exc()))
+        else:
+            self.signals.result.emit(result)
+        finally:
+            self.signals.finished.emit()

core/scan.py

@@ -345,8 +345,7 @@ def disasm(self, sec_bytes: bytes, base_va: int, section_offset: int, syscall_ma
                 for ins in result.get("context", [])[-12:]:
                     mnemonic = str(ins['mnemonic'])
                     arrow = ">>> " if mnemonic.startswith("svc") else "    "
-                    info(
-                        f"  {arrow}0x{ins['address']:x}  {mnemonic} {ins['op_str']}", "debug")
+                    debug(f"  {arrow}0x{ins['address']:x}  {mnemonic} {ins['op_str']}")
                 debug("")
             i += 4
         return results

img/spinner.gif

@@ -1,11 +1,21 @@
+import json
+
 from PySide6.QtWidgets import (
-    QWidget, QVBoxLayout, QPushButton, QLabel, QTreeWidget, QTreeWidgetItem, QFileDialog, QHBoxLayout
+    QWidget, QVBoxLayout, QPushButton, QLabel,
+    QTreeWidget, QTreeWidgetItem, QFileDialog, QHBoxLayout
 )
 
+from core.concurrency.subprocess import Subprocess
 from core.scan import MachOScanner
 from core.utils import is_macho_file
 
 
+def run_scan_process(file_path: str, verbose: bool = True):
+    scanner = MachOScanner()
+    results = scanner.analyze(file_path, verbose=verbose, out_path=None)
+    return json.dumps(results)
+
+
 class ScannerTab(QWidget):
     def __init__(self):
         super().__init__()
@@ -25,6 +35,8 @@ def __init__(self):
         self.tree.setColumnWidth(0, 220)
         self.tree.setAlternatingRowColors(True)
 
+        # todo add spinner
+
         self.scan_button = QPushButton("Run Scan")
         self.scan_button.clicked.connect(self._on_scan_clicked)
         self.scan_button.setEnabled(False)
@@ -36,7 +48,11 @@ def __init__(self):
         self.setAcceptDrops(True)
 
         self.file_path = None
-        self.scanner = MachOScanner()
+        self.worker = Subprocess()
+
+    def _toggle_button_states(self, on):
+        self.scan_button.setEnabled(on)
+        self.open_button.setEnabled(on)
 
     def dragEnterEvent(self, event):
         if event.mimeData().hasUrls():
@@ -49,6 +65,7 @@ def dragEnterEvent(self, event):
             event.ignore()
 
     def dropEvent(self, event):
+        # todo forbid drop when already scanning
         if event.mimeData().hasUrls():
             file_path = event.mimeData().urls()[0].toLocalFile()
             if file_path and is_macho_file(file_path):
@@ -80,22 +97,26 @@ def _run_scan(self):
         if not self.file_path:
             return
 
-        self.scan_button.setEnabled(False)
+        self._toggle_button_states(False)
         self.tree.clear()
         status_root = QTreeWidgetItem(["Status", "Scanning..."])
         self.tree.addTopLevelItem(status_root)
 
-        # TODO: This is a blocking call, UI will freeze here for large files!
-        results = self.scanner.analyze(self.file_path, verbose=True)
-
-        print(results)
+        self.worker.submit(
+            run_scan_process,
+            self.file_path,
+            True,
+            on_done=self._display_results,
+            on_error=self._display_error
+        )
 
+    def _display_results(self, jsonResults):
         self.tree.clear()
         root = QTreeWidgetItem(["Scan Results", self.file_path])
         self.tree.addTopLevelItem(root)
 
-        for i, result in enumerate(results, 1):
-            result_item = QTreeWidgetItem([f"Binary"])
+        for i, result in enumerate(json.loads(jsonResults), 1):
+            result_item = QTreeWidgetItem(["Binary"])
             root.addChild(result_item)
 
             for category, data in result.items():
@@ -109,15 +130,6 @@ def _run_scan(self):
                             f"section: {s.get('section', '')}, offset: {s.get('offset', '')}"
                         ])
                         cat_item.addChild(syscall_item)
-
-                        # context = s.get("context", [])
-                        # for ins in context[-12:]:  # last 12 instructions
-                        #    ctx_item = QTreeWidgetItem([
-                        #        f"0x{ins['address']:x} {ins['mnemonic']}",
-                        #        ins.get("op_str", "")
-                        #    ])
-                        #    syscall_item.addChild(ctx_item)
-
                 elif isinstance(data, dict):
                     for k, v in data.items():
                         sub_item = QTreeWidgetItem([str(k), str(v)])
@@ -131,4 +143,13 @@ def _run_scan(self):
 
         root.setExpanded(True)
         self.tree.expandAll()
-        self.scan_button.setEnabled(True)
+        self._toggle_button_states(True)
+
+    def _display_error(self, error):
+        self.tree.clear()
+        self.tree.addTopLevelItem(QTreeWidgetItem(["Error", str(error)]))
+        self._toggle_button_states(True)
+
+    def closeEvent(self, event):
+        self.worker.shutdown(wait=False, cancel_futures=True)
+        super().closeEvent(event)