From d212d2e8891509584c84f204dd7d22cce8f07bb4 Mon Sep 17 00:00:00 2001
From: Alexandre Iooss <erdnaxe@crans.org>
Date: Wed, 27 Sep 2023 10:54:47 +0200
Subject: [PATCH] Stub by address, not by function name

---
 rainbow/rainbow.py | 40 +++++++++++++++++++++++++++++++---------
 tests/test_hook.py |  4 ++--
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/rainbow/rainbow.py b/rainbow/rainbow.py
index 9b72fd0..6dfdb9b 100644
--- a/rainbow/rainbow.py
+++ b/rainbow/rainbow.py
@@ -418,15 +418,24 @@ def hook_prolog(self, name, fn):
         """
         if not self.allow_stubs:
             raise ValueError("Cannot use stubs, allow_stubs is False.")
-        if name not in self.functions.keys():
-            raise IndexError(f"'{name}' could not be found.")
 
         def to_hook(x):
             if fn is not None:
                 fn(x)
             return False
 
-        self.stubbed_functions[name] = to_hook
+        if isinstance(name, str):
+            # Stub all function addresses matching this name
+            addrs = [a for a, n in self.function_names.items() if n == name]
+            if not addrs:
+                raise IndexError(f"'{name}' could not be found.")
+            for addr in addrs:
+                self.stubbed_functions[addr] = to_hook
+        elif isinstance(name, int):
+            # Name is an address
+            self.stubbed_functions[name] = to_hook
+        else:
+            raise TypeError("name should be function name or address")
 
     def hook_bypass(self, name, fn=None):
         """
@@ -435,15 +444,24 @@ def hook_bypass(self, name, fn=None):
         """
         if not self.allow_stubs:
             raise ValueError("Cannot use stubs, allow_stubs is False.")
-        if name not in self.functions.keys():
-            raise IndexError(f"'{name}' could not be found.")
 
         def to_hook(x):
             if fn is not None:
                 fn(x)
             return True
 
-        self.stubbed_functions[name] = to_hook
+        if isinstance(name, str):
+            # Stub all function addresses matching this name
+            addrs = [a for a, n in self.function_names.items() if n == name]
+            if not addrs:
+                raise IndexError(f"'{name}' could not be found.")
+            for addr in addrs:
+                self.stubbed_functions[addr] = to_hook
+        elif isinstance(name, int):
+            # Name is an address
+            self.stubbed_functions[name] = to_hook
+        else:
+            raise TypeError("name should be function name or address")
 
     def remove_hook(self, name):
         """Remove the hook."""
@@ -463,16 +481,20 @@ def _block_hook(self, _uci, address: int, _size, _):
         address+function is redefined in the user's python script and if so,
         calls that instead.
         """
+        # Print function calls
         if address in self.function_names and (self.allow_stubs or self.print_config & Print.Functions):
             # Handle the function call printing
             f = self.function_names[address]
             if self.print_config & Print.Functions:
                 print(f"{color('MAGENTA', f)}(...) @ 0x{address:x}")
 
-            # Handle the stubs
-            if f in self.stubbed_functions:
-                r = self.stubbed_functions[f](self)
+        # If stub is enabled and set at this address, run it
+        if self.allow_stubs:
+            stub_func = self.stubbed_functions.get(address)
+            if stub_func is not None:
+                r = stub_func(self)
                 if r:
+                    # If stub returns True, then make the function return early
                     self.return_force()
 
     def _mem_hook(self, uci, access, address, size, value, _):
diff --git a/tests/test_hook.py b/tests/test_hook.py
index c460be7..441fc61 100644
--- a/tests/test_hook.py
+++ b/tests/test_hook.py
@@ -46,6 +46,6 @@ def test_remove_hooks():
     emu.setup()
 
     emu.hook_bypass("strtol")
-    assert "strtol" in emu.stubbed_functions
+    assert 0x202fa0 in emu.stubbed_functions
     emu.remove_hooks()
-    assert "strtol" not in emu.stubbed_functions
+    assert 0x202fa0 not in emu.stubbed_functions