Implement BugHouse components plugin. (#179)

Also with the following improvements:

- Plugins now support headless mode.
- URL scheme handling now only requires implementing a single handler
(instead of three).

angrmanagement/__main__.py

@@ -200,6 +200,10 @@ def main():
             run_daemon_process()
             time.sleep(1)
 
+        # initialize plugins
+        from .plugins import PluginManager
+        PluginManager(None).discover_and_initialize_plugins()
+
         action = handle_url(args.url, act=False)
         action.act(daemon_conn())
         return

angrmanagement/daemon/client.py

@@ -65,12 +65,28 @@ def exposed_openbitmap(self, bitmap_path, base):
                 _l.critical("TracePlugin is probably not installed.")
                 # TODO: Open a message box
 
+    def exposed_custom_binary_aware_action(self, action, kwargs):
+        kwargs_copy = dict(kwargs.items())  # copy it to local
+        DaemonClient.invoke(action, kwargs_copy)
+
 
 class DaemonClientCls:
     """
     Implements logic that the client needs to talk to the daemon service.
     """
 
+    def __init__(self, custom_handlers=None):
+        self.custom_handlers = {}
+
+    def register_handler(self, action:str, handler):
+        self.custom_handlers[action] = handler
+
+    def invoke(self, action, kwargs):
+        if action not in self.custom_handlers:
+            _l.critical("Unregistered URL action \"%s\"." % action)
+            return
+        self.custom_handlers[action](kwargs)
+
     @property
     def conn(self):
         return GlobalInfo.daemon_conn

angrmanagement/daemon/server.py

@@ -91,6 +91,10 @@ def exposed_openbitmap(self, bitmap_path, base, md5, sha256):
     def exposed_exit(self):
         pass
 
+    def exposed_custom_binary_aware_action(self, md5, sha256, action, kwargs):
+        conn = self._get_conn(md5, sha256)
+        conn.root.custom_binary_aware_action(action, kwargs)
+
 
 def monitor_thread(server):
     """
@@ -128,7 +132,13 @@ def start_daemon(port=DEFAULT_PORT):
     except SingleInstanceException:
         return
 
-    server = ThreadedServer(ManagementService, port=port)
+    # load plugins in headless mode
+    from ..plugins import PluginManager
+    GlobalInfo.headless_plugin_manager = PluginManager(None)
+    GlobalInfo.headless_plugin_manager.discover_and_initialize_plugins()
+
+    # start the server
+    server = ThreadedServer(ManagementService, port=port, protocol_config={'allow_public_attrs': True})
     threading.Thread(target=monitor_thread, args=(server, ), daemon=True).start()
     server.start()
 
@@ -163,5 +173,6 @@ def daemon_conn(port=DEFAULT_PORT, service=None):
     kwargs = { }
     if service is not None:
         kwargs['service'] = service
+    kwargs['config'] = {'allow_public_attrs': True}
     conn = rpyc.connect("localhost", port, **kwargs)
     return conn

angrmanagement/daemon/url_handler.py

@@ -142,6 +142,40 @@ def _from_params(cls, params):
         )
 
 
+class UrlActionBinaryAware(UrlActionBase):
+    """
+    The base class of all binary-aware URl actions.
+    """
+    def __init__(self, md5=None, sha256=None, action=None, kwargs=None):
+        super().__init__(md5, sha256)
+        self.action = action
+        self.kwargs = kwargs
+
+        if not self.md5 and not self.sha256:
+            raise TypeError("You must provide either MD5 or SHA256 of the target binary.")
+        if not self.action:
+            raise TypeError("You must provide action.")
+
+    def act(self, daemon_conn=None):
+        daemon_conn.root.custom_binary_aware_action(
+            self.md5, self.sha256, self.action, self.kwargs
+        )
+
+    @classmethod
+    def _from_params(cls, params):
+        sha256 = cls._one_param(params, 'sha256')
+        md5 = cls._one_param(params, 'md5')
+        action = cls._one_param(params, 'action')
+        kwargs = {}
+        for k, v in params.items():
+            if k not in {'sha256', 'md5', 'action'}:
+                if isinstance(v, (list, tuple)):
+                    kwargs[k] = v[0]
+                else:
+                    kwargs[k] = v
+        return cls(md5=md5, sha256=sha256, action=action, kwargs=kwargs)
+
+
 _ACT2CLS = {
     'open': UrlActionOpen,
     'jumpto': UrlActionJumpTo,
@@ -158,3 +192,7 @@ def handle_url(url, act=True):
     if act:
         action.act()
     return action
+
+
+def register_url_action(action: str, action_handler: UrlActionBase):
+    _ACT2CLS[action] = action_handler

angrmanagement/data/instance.py

@@ -62,7 +62,7 @@ def __init__(self, project=None):
     #
 
     @property
-    def project(self):
+    def project(self) -> Optional[angr.Project]:
         return self._project_container.am_obj
 
     @project.setter

angrmanagement/data/jobs/variable_recovery.py

@@ -1,10 +1,11 @@
-
+from typing import TYPE_CHECKING
 import time
 
-from angr.analyses import CallingConventionAnalysis
-
 from .job import Job
 
+if TYPE_CHECKING:
+    from ..instance import Instance
+
 
 class VariableRecoveryJob(Job):
     """
@@ -16,10 +17,11 @@ def __init__(self, on_finish=None):
 
         self._last_progress_callback_triggered = None
 
-    def run(self, inst):
+    def run(self, inst: 'Instance'):
         inst.project.analyses.CompleteCallingConventions(
             recover_variables=True,
             low_priority=True,
+            cfg=inst.cfg,
             progress_callback=self._progress_callback,
         )
 

angrmanagement/logic/__init__.py

@@ -4,3 +4,4 @@ class GlobalInfo:
     main_window = None
     daemon_inst = None
     daemon_conn = None
+    headless_plugin_manager = None

angrmanagement/plugins/base_plugin.py

@@ -12,11 +12,14 @@
 
 
 class BasePlugin:
+    REQUIRE_WORKSPACE = True
     __i_hold_this_abstraction_token = True
 
-
     def __init__(self, workspace):
-        self.workspace = workspace
+
+        from angrmanagement.ui.workspace import Workspace
+
+        self.workspace: Optional[Workspace] = workspace
         _l.info("Loaded plugin {}".format(self.__class__.__name__))
 
         # valid things that we want you do be able to do in __init__:
@@ -60,16 +63,19 @@ def handle_click_block(self, qblock, event: QGraphicsSceneMouseEvent):
 
     # iterable of tuples (icon, tooltip)
     TOOLBAR_BUTTONS = []  # type: List[Tuple[QIcon, str]]
+
     def handle_click_toolbar(self, idx):
         pass
 
     # Iterable of button texts
     MENU_BUTTONS = []  # type: List[str]
+
     def handle_click_menu(self, idx):
         pass
 
     # Iterable of column names
     FUNC_COLUMNS = []  # type: List[str]
+
     def extract_func_column(self, func, idx) -> Tuple[Any, str]:
         # should return a tuple of the sortable column data and the rendered string
         return 0, ''
@@ -85,3 +91,9 @@ def build_context_menu_block(self, item) -> Iterator[Union[None, Tuple[str, Call
         Use None to insert a MenuSeparator(). The tuples are: (menu entry text, callback)
         """
         return []
+
+    # Iterable of URL actions
+    URL_ACTIONS: List[str] = []
+
+    def handle_url_action(self, action, kwargs):
+        pass

angrmanagement/plugins/bughouse/__init__.py

@@ -0,0 +1 @@
+from .components import ComponentsPlugin

angrmanagement/plugins/bughouse/components.py

@@ -0,0 +1,90 @@
+from typing import Optional
+
+from angrmanagement.logic.threads import is_gui_thread, gui_thread_schedule_async
+from angrmanagement.plugins.base_plugin import BasePlugin
+
+from .ui import LoadComponentsDialog, ComponentsView
+
+
+class ComponentsPlugin(BasePlugin):
+    REQUIRE_WORKSPACE = False
+
+    """
+    Implement a component viewer that takes JSON messages on function clustering from the server side and visualizes
+    the clusters in a treeview.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if self.workspace is not None:
+            self.view: ComponentsView = ComponentsView(self.workspace, "left")
+
+            # register a new view
+            self.workspace.view_manager.add_view(self.view, self.view.caption, self.view.category)
+
+    def teardown(self):
+        pass
+
+    #
+    # Menus
+    #
+
+    MENU_BUTTONS = [
+        'Load components...',
+        'Reset components',
+    ]
+    LOAD_COMPONENTS = 0
+    RESET_COMPONENTS = 1
+
+    def handle_click_menu(self, idx):
+
+        if idx < 0 or idx >= len(self.MENU_BUTTONS):
+            return
+
+        if self.workspace.instance.project is None:
+            return
+
+        mapping = {
+            self.LOAD_COMPONENTS: self.load_components,
+            self.RESET_COMPONENTS: self.reset_components,
+        }
+        mapping.get(idx)()
+
+    def load_components(self, url: Optional[str]=None):
+        """
+        Open a new dialog and take a JSON URL or a file path. Then load components from that URL.
+        """
+        dialog = LoadComponentsDialog(workspace=self.workspace, url=url)
+        dialog.exec_()
+        if dialog.tree is not None:
+            self.view.load(dialog.tree)
+
+    def reset_components(self):
+        """
+        Clear existing components information.
+        """
+        self.view.reset()
+
+    #
+    # URLs
+    #
+
+    # register actions
+    URL_ACTIONS = [
+        'bughouse_component',
+    ]
+
+    def handle_url_action(self, action, kwargs):
+        mapping = {
+            'bughouse_component': self.handle_url_action_bughouse_component,
+        }
+
+        func = mapping.get(action)
+        if is_gui_thread():
+            func(**kwargs)
+        else:
+            gui_thread_schedule_async(func, kwargs=kwargs)
+
+    def handle_url_action_bughouse_component(self, url=None):
+        self.load_components(url)

angrmanagement/plugins/bughouse/data/__init__.py

@@ -0,0 +1 @@
+from .component_tree import ComponentTreeNode, ComponentTree, ComponentFunction

angrmanagement/plugins/bughouse/data/component_tree.py

@@ -0,0 +1,43 @@
+from typing import List, Optional
+
+
+class ComponentFunction:
+
+    __slots__ = ('mapped_base', 'virtual_addr', 'symbol_name', )
+
+    def __init__(self, mapped_base: int, virtual_addr: int, symbol_name: Optional[str]=None):
+        self.mapped_base = mapped_base
+        self.virtual_addr = virtual_addr
+        self.symbol_name = symbol_name
+
+    def __eq__(self, other):
+        return isinstance(other, ComponentFunction) and \
+                self.mapped_base == other.mapped_base and \
+                self.virtual_addr == other.virtual_addr
+
+    def __hash__(self):
+        return hash((ComponentFunction, self.mapped_base, self.virtual_addr))
+
+
+class ComponentTreeNode:
+    def __init__(self, name=None):
+        self.name = name
+        self.components: List['ComponentTreeNode'] = [ ]
+        self.functions: List[ComponentFunction] = [ ]
+
+    def __eq__(self, other):
+        return isinstance(other, ComponentTreeNode) \
+                and self.components == other.components \
+                and set(self.functions) == set(other.functions)
+
+    def __hash__(self):
+        return hash((ComponentTreeNode,
+                     hash(tuple(self.components)),
+                     hash(tuple(sorted((f.mapped_base + f.virtual_addr) for f in self.functions))),
+                     )
+                    )
+
+
+class ComponentTree:
+    def __init__(self, root: Optional[ComponentTreeNode]=None):
+        self.root = root

angrmanagement/plugins/bughouse/ui/__init__.py

@@ -0,0 +1,2 @@
+from .load_components_dialog import LoadComponentsDialog
+from .components_treeview import ComponentsView