[Binja] Fix custom hook persistence and tests (#99)

Author: lordidiot

plugins/binaryninja/mui/hook_manager.py

@@ -121,7 +121,7 @@ def load_existing_hooks(self) -> None:
             json.loads(settings.get_string(f"{BINJA_HOOK_SETTINGS_PREFIX}avoid", bv))
         )
         bv.session_data.mui_custom_hooks = {
-            key: item
+            CustomHookIdentity.from_name(key): item
             for key, item in json.loads(
                 settings.get_string(f"{BINJA_HOOK_SETTINGS_PREFIX}custom", bv)
             ).items()
@@ -168,9 +168,12 @@ def save_hooks(self) -> None:
             scope=SettingsScope.SettingsResourceScope,
         )
 
+        custom_hooks = {
+            key.to_name(): item for key, item in bv.session_data.mui_custom_hooks.items()
+        }
         settings.set_string(
             f"{BINJA_HOOK_SETTINGS_PREFIX}custom",
-            json.dumps(bv.session_data.mui_custom_hooks),
+            json.dumps(custom_hooks),
             view=bv,
             scope=SettingsScope.SettingsResourceScope,
         )

plugins/binaryninja/tests/binaries/hello_world_basic_hooks.bndb

@@ -5,7 +5,7 @@
 
 from manticore.native import Manticore
 from mui.native_plugin import RebaseHooksPlugin
-from mui.hook_manager import NativeHookManager
+from mui.hook_manager import NativeHookManager, CustomHookIdentity
 
 
 class FakeHookManager:
@@ -14,7 +14,7 @@ def __init__(
         filename: str,
         find_hooks: Set[int] = set(),
         avoid_hooks: Set[int] = set(),
-        custom_hooks: Dict[str, str] = {},
+        custom_hooks: Dict[CustomHookIdentity, str] = {},
         global_hooks: Dict[str, str] = {},
     ):
         self.bv = MagicMock()
@@ -30,7 +30,7 @@ def list_find_hooks(self) -> Set[int]:
     def list_avoid_hooks(self) -> Set[int]:
         return self.avoid_hooks
 
-    def list_custom_hooks(self) -> Dict[str, str]:
+    def list_custom_hooks(self) -> Dict[CustomHookIdentity, str]:
         return self.custom_hooks
 
     def list_global_hooks(self) -> Dict[str, str]:
@@ -104,7 +104,9 @@ def test_rebase_custom(self) -> None:
 
         mgr = cast(
             NativeHookManager,
-            FakeHookManager(self.LIB_PATH, custom_hooks={f"{self.FOO}_00": custom_code}),
+            FakeHookManager(
+                self.LIB_PATH, custom_hooks={CustomHookIdentity(self.FOO, 0): custom_code}
+            ),
         )
         m.register_plugin(RebaseHooksPlugin(mgr, self.find_f, self.avoid_f))
         m.run()

plugins/binaryninja/tests/test_native_plugin.py

@@ -1,6 +1,6 @@
 import os
 import unittest
-from mui.hook_manager import NativeHookManager
+from mui.hook_manager import NativeHookManager, CustomHookIdentity
 from binaryninja import open_view
 
 
@@ -13,7 +13,9 @@ def test_hook_loading(self) -> None:
             mgr.load_existing_hooks()
             self.assertEqual(mgr.list_find_hooks(), set([0x401152]))
             self.assertEqual(mgr.list_avoid_hooks(), set([0x401153]))
-            self.assertEqual(mgr.list_custom_hooks(), {"0x40114D": "# custom hook code"})
+            self.assertEqual(
+                mgr.list_custom_hooks(), {CustomHookIdentity(0x40114D, 0): "# custom hook code"}
+            )
             self.assertEqual(mgr.list_global_hooks(), {"global_00": "# global hook code"})