Add a signatures view.

angrmanagement/data/instance.py

@@ -11,6 +11,7 @@
 from cle import SymbolType
 
 from angrmanagement.data.breakpoint import Breakpoint, BreakpointManager, BreakpointType
+from angrmanagement.data.signatures import SignatureManager
 from angrmanagement.data.trace import Trace
 from angrmanagement.errors import ContainerAlreadyRegisteredError
 from angrmanagement.logic.debugger import DebuggerListManager, DebuggerManager
@@ -65,6 +66,7 @@ def __init__(self) -> None:
         self.register_container("active_view_state", lambda: None, "ViewState", "Currently focused view state")
 
         self.breakpoint_mgr = BreakpointManager()
+        self.signature_mgr = SignatureManager(self)
         self.debugger_list_mgr = DebuggerListManager()
         self.debugger_mgr = DebuggerManager(self.debugger_list_mgr)
 

angrmanagement/data/signatures.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING
+
+import angr
+import angr.flirt
+from angr.flirt import FLIRT_SIGNATURES_BY_ARCH, FlirtSignature
+
+from angrmanagement.config import Conf
+from angrmanagement.utils.env import app_root, is_pyinstaller
+
+from .object_container import ObjectContainer
+
+if TYPE_CHECKING:
+    from .instance import Instance
+
+_l = logging.getLogger(__name__)
+
+#
+# FLIRT Signatures
+#
+
+
+def init_flirt_signatures() -> None:
+    if Conf.flirt_signatures_root:
+        # if it's a relative path, it's relative to the angr-management package
+        if os.path.isabs(Conf.flirt_signatures_root):
+            flirt_signatures_root = Conf.flirt_signatures_root
+        else:
+            if is_pyinstaller():
+                flirt_signatures_root = os.path.join(app_root(), Conf.flirt_signatures_root)
+            else:
+                # when running as a Python package, we should use the git submodule, which is on the same level
+                # with (instead of inside) the angrmanagement module directory.
+                flirt_signatures_root = os.path.join(app_root(), "..", Conf.flirt_signatures_root)
+        flirt_signatures_root = os.path.normpath(flirt_signatures_root)
+        _l.info("Loading FLIRT signatures from %s.", flirt_signatures_root)
+        angr.flirt.load_signatures(flirt_signatures_root)
+
+
+class SignatureManager:
+    """
+    Manager of function signatures.
+    """
+
+    def __init__(self, instance: Instance) -> None:
+        self.signatures: ObjectContainer = ObjectContainer([], "List of function signatures")
+        self.dryrun_results: dict[str, dict[int, str]] = {}
+        self.instance = instance
+
+    def sync_from_angr(self):
+        for _arch, sigs in FLIRT_SIGNATURES_BY_ARCH.items():
+            for sig in sigs:
+                self.add_signature(sig)
+
+    def clear(self) -> None:
+        self.signatures.clear()
+        self.signatures.am_event()
+
+    def add_signature(self, sig: FlirtSignature) -> None:
+        self.signatures.append(sig)
+        self.signatures.am_event(added=sig)
+
+    def remove_signature(self, sig: FlirtSignature) -> None:
+        self.signatures.remove(sig)
+        self.signatures.am_event(removed=sig)
+
+    def apply_signatures(self, sigs: list[FlirtSignature], dry_run: bool = True):
+        for sig in sigs:
+            fl = self.instance.project.analyses.Flirt(sig.sig_path, dry_run=dry_run)
+            if dry_run:
+                self.dryrun_results[sig.sig_path] = next(iter(fl.matched_suggestions.values()))[1]
+
+    def get_match_count(self, sig: FlirtSignature) -> int | None:
+        if sig.sig_path in self.dryrun_results:
+            return len(self.dryrun_results[sig.sig_path])
+        return None

angrmanagement/ui/main_window.py

@@ -8,8 +8,6 @@
 from functools import partial
 from typing import TYPE_CHECKING
 
-import angr
-import angr.flirt
 import PySide6QtAds as QtAds
 from angr.angrdb import AngrDB
 from PySide6.QtCore import QEvent, QObject, QSize, Qt, QUrl
@@ -28,14 +26,14 @@
 from angrmanagement.data.jobs import DependencyAnalysisJob
 from angrmanagement.data.jobs.loading import LoadAngrDBJob, LoadBinaryJob
 from angrmanagement.data.library_docs import LibraryDocs
+from angrmanagement.data.signatures import init_flirt_signatures
 from angrmanagement.errors import InvalidURLError, UnexpectedStatusCodeError
 from angrmanagement.logic import GlobalInfo
 from angrmanagement.logic.commands import BasicCommand
 from angrmanagement.logic.threads import ExecuteCodeEvent
 from angrmanagement.ui.dialogs.progress_dialog import ProgressDialog
 from angrmanagement.ui.views import DisassemblyView
 from angrmanagement.ui.widgets.qam_status_bar import QAmStatusBar
-from angrmanagement.utils.env import app_root, is_pyinstaller
 from angrmanagement.utils.io import download_url, isurl
 
 from .dialogs.about import LoadAboutDialog
@@ -181,7 +179,7 @@ def __init__(
         self.workspace.plugins.on_workspace_initialized(self)
 
         self._init_shortcuts()
-        self._init_flirt_signatures()
+        init_flirt_signatures()
 
         self._run_daemon(use_daemon=use_daemon)
 
@@ -379,26 +377,6 @@ def init_shortcuts_on_dock(self, dock_widget) -> None:
     def _init_plugins(self) -> None:
         self.workspace.plugins.discover_and_initialize_plugins()
 
-    #
-    # FLIRT Signatures
-    #
-
-    def _init_flirt_signatures(self) -> None:
-        if Conf.flirt_signatures_root:
-            # if it's a relative path, it's relative to the angr-management package
-            if os.path.isabs(Conf.flirt_signatures_root):
-                flirt_signatures_root = Conf.flirt_signatures_root
-            else:
-                if is_pyinstaller():
-                    flirt_signatures_root = os.path.join(app_root(), Conf.flirt_signatures_root)
-                else:
-                    # when running as a Python package, we should use the git submodule, which is on the same level
-                    # with (instead of inside) the angrmanagement module directory.
-                    flirt_signatures_root = os.path.join(app_root(), "..", Conf.flirt_signatures_root)
-            flirt_signatures_root = os.path.normpath(flirt_signatures_root)
-            _l.info("Loading FLIRT signatures from %s.", flirt_signatures_root)
-            angr.flirt.load_signatures(flirt_signatures_root)
-
     #
     # Library docs
     #
@@ -512,6 +490,7 @@ def _register_commands(self) -> None:
                     ("View: Disassembly (Linear)", self.workspace.show_linear_disassembly_view),
                     ("View: Functions", self.workspace.show_functions_view),
                     ("View: Hex", self.workspace.show_hex_view),
+                    ("View: Function Signatures", self.workspace.show_signatures_view),
                     ("View: Jobs", self.workspace.show_jobs_view),
                     ("View: Log", self.workspace.show_log_view),
                     ("View: New Disassembly (Graph)", self.workspace.create_and_show_graph_disassembly_view),

angrmanagement/ui/menus/view_menu.py

@@ -128,8 +128,9 @@ def __init__(self, main_window: MainWindow) -> None:
                 MenuEntry("&Patches", main_window.workspace.show_patches_view, icon=icon("patches-view")),
                 MenuEntry("&Types", main_window.workspace.show_types_view, icon=icon("types-view")),
                 MenuEntry("&Functions", main_window.workspace.show_functions_view, icon=icon("functions-view")),
+                MenuEntry("Function Si&gnatures", main_window.workspace.show_signatures_view),
                 MenuEntry("&Traces", main_window.workspace.show_traces_view, icon=icon("traces-view")),
-                MenuEntry("&Trace Map", main_window.workspace.show_trace_map_view),
+                MenuEntry("Trace &Map", main_window.workspace.show_trace_map_view),
                 MenuSeparator(),
                 MenuEntry("Symbolic &Execution", main_window.workspace.show_symexec_view),
                 MenuEntry("S&ymbolic States", main_window.workspace.show_states_view),

angrmanagement/ui/views/__init__.py

@@ -14,6 +14,7 @@
 from .patches_view import PatchesView
 from .proximity_view import ProximityView
 from .registers_view import RegistersView
+from .signatures_view import SignaturesView
 from .stack_view import StackView
 from .states_view import StatesView
 from .strings_view import StringsView
@@ -42,6 +43,7 @@
     "StatesView",
     "StringsView",
     "SymexecView",
+    "SignaturesView",
     "TraceMapView",
     "TracesView",
     "TypesView",

angrmanagement/ui/views/signatures_view.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from PySide6.QtCore import QAbstractTableModel, QSize, Qt
+from PySide6.QtWidgets import QAbstractItemView, QHeaderView, QMenu, QTableView, QVBoxLayout
+
+from .view import InstanceView
+
+if TYPE_CHECKING:
+
+    import PySide6
+    from angr.flirt import FlirtSignature
+
+    from angrmanagement.data.instance import Instance
+    from angrmanagement.data.signatures import SignatureManager
+    from angrmanagement.ui.workspace import Workspace
+
+
+class QSignatureTableModel(QAbstractTableModel):
+    """
+    Signature table model.
+    """
+
+    Headers = [
+        "Type",
+        "Name",
+        "Architecture",
+        "Platform",
+        "Compiler",
+        "OS Name",
+        "Matches",
+    ]
+    COL_TYPE = 0
+    COL_NAME = 1
+    COL_ARCH = 2
+    COL_PLATFORM = 3
+    COL_COMPILER = 4
+    COL_OS_NAME = 5
+    COL_MATCHES = 6
+
+    def __init__(self, signature_mgr: SignatureManager) -> None:
+        super().__init__()
+        self.signature_mgr = signature_mgr
+        self.signature_mgr.signatures.am_subscribe(self._on_signatures_updated)
+
+    def _on_signatures_updated(self, **kwargs) -> None:  # pylint:disable=unused-argument
+        self.beginResetModel()
+        self.endResetModel()
+
+    def rowCount(self, parent: PySide6.QtCore.QModelIndex = ...) -> int:  # pylint:disable=unused-argument
+        return len(self.signature_mgr.signatures)
+
+    def columnCount(self, parent: PySide6.QtCore.QModelIndex = ...) -> int:  # pylint:disable=unused-argument
+        return len(self.Headers)
+
+    def headerData(
+        self, section: int, orientation: PySide6.QtCore.Qt.Orientation, role: int = ...
+    ) -> Any:  # pylint:disable=unused-argument
+        if role != Qt.ItemDataRole.DisplayRole:
+            return None
+        if section < len(self.Headers):
+            return self.Headers[section]
+        return None
+
+    def data(self, index: PySide6.QtCore.QModelIndex, role: int = ...) -> Any:
+        if not index.isValid():
+            return None
+        row = index.row()
+        if row >= len(self.signature_mgr.signatures):
+            return None
+        col = index.column()
+        if role == Qt.ItemDataRole.DisplayRole:
+            return self._get_column_text(self.signature_mgr.signatures[row], col)
+        else:
+            return None
+
+    def _get_column_text(self, sig: FlirtSignature, column: int) -> str:
+        if column == self.COL_TYPE:
+            return "FLIRT"
+        elif column == self.COL_NAME:
+            return sig.sig_name
+        elif column == self.COL_ARCH:
+            return sig.arch
+        elif column == self.COL_PLATFORM:
+            return sig.platform
+        elif column == self.COL_COMPILER:
+            return sig.compiler
+        elif column == self.COL_OS_NAME:
+            return sig.os_name
+        elif column == self.COL_MATCHES:
+            m = self.signature_mgr.get_match_count(sig)
+            return str(m) if m is not None else "N/A"
+        else:
+            raise AssertionError
+
+
+class QSignatureTableWidget(QTableView):
+    """
+    Signature table widget.
+    """
+
+    def __init__(self, signature_mgr, workspace: Workspace) -> None:
+        super().__init__()
+        self.workspace = workspace
+        self.signature_mgr = signature_mgr
+
+        hheader = self.horizontalHeader()
+        hheader.setVisible(True)
+
+        vheader = self.verticalHeader()
+        vheader.setVisible(False)
+        vheader.setDefaultSectionSize(20)
+
+        self.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows)
+        self.setHorizontalScrollMode(QAbstractItemView.ScrollMode.ScrollPerPixel)
+
+        self.model: QSignatureTableModel = QSignatureTableModel(self.signature_mgr)
+        self.setModel(self.model)
+
+        for col in range(len(QSignatureTableModel.Headers)):
+            hheader.setSectionResizeMode(col, QHeaderView.ResizeMode.ResizeToContents)
+        hheader.setStretchLastSection(True)
+        self.doubleClicked.connect(self._on_cell_double_click)
+
+    #
+    # Events
+    #
+
+    def closeEvent(self, event) -> None:
+        self.model.shutdown()
+        super().closeEvent(event)
+
+    def contextMenuEvent(self, event) -> None:
+        selected_rows = {i.row() for i in self.selectedIndexes()}
+        sigs = [self.signature_mgr.signatures[r] for r in selected_rows]
+        menu = QMenu("", self)
+        if len(sigs):
+            menu.addAction("Try applying signature(s)", lambda: self.signature_mgr.apply_signatures(sigs, dry_run=True))
+            menu.addAction("Apply signature(s)", lambda: self.signature_mgr.apply_signatures(sigs, dry_run=False))
+            menu.addSeparator()
+        # menu.addAction("Load signature file", self.signature_mgr.load_signature)
+        menu.exec_(event.globalPos())
+
+    def _on_cell_double_click(self, index) -> None:
+        return
+
+
+class SignaturesView(InstanceView):
+    """
+    Signatures view that displays all loaded (FLIRT) signatures and their statuses.
+    """
+
+    def __init__(self, workspace: Workspace, default_docking_position: str, instance: Instance) -> None:
+        super().__init__("signatures", workspace, default_docking_position, instance)
+        self.base_caption = "Function Signatures"
+        self._tbl_widget: QSignatureTableWidget | None = None
+        self._init_widgets()
+        self.reload()
+
+    def reload(self) -> None:
+        self.instance.signature_mgr.sync_from_angr()
+
+    def minimumSizeHint(self):
+        return QSize(200, 200)
+
+    def _init_widgets(self) -> None:
+        vlayout = QVBoxLayout()
+        self._tbl_widget = QSignatureTableWidget(self.instance.signature_mgr, self.workspace)
+        vlayout.addWidget(self._tbl_widget)
+        self.setLayout(vlayout)

angrmanagement/ui/workspace.py

@@ -67,6 +67,7 @@
     PatchesView,
     ProximityView,
     RegistersView,
+    SignaturesView,
     StackView,
     StatesView,
     StringsView,
@@ -821,6 +822,9 @@ def show_functions_view(self) -> None:
     def show_traces_view(self) -> None:
         self.show_view("traces", TracesView)
 
+    def show_signatures_view(self) -> None:
+        self.show_view("signatures", SignaturesView)
+
     def show_trace_map_view(self) -> None:
         self.show_view("tracemap", TraceMapView, position="top")