add create_capture_hook test

Author: StarrFox

memobj/allocation.py

@@ -78,7 +78,7 @@ def __exit__(self, *_):
     def closed(self) -> bool:
         return self._is_closed
 
-    def allocate(self, size: int) -> Allocation:
+    def allocate(self, size: int, *, preferred_start: int | None = None) -> Allocation:
         """
         Allocates a block of memory for the process.
 
@@ -89,11 +89,12 @@ def allocate(self, size: int) -> Allocation:
 
         Args:
             size (int): The size of the memory block to allocate in bytes.
+            preferred_start: The preferred start address of the allocation
 
         Returns:
             Allocation: An object representing the allocated memory block.
         """
-        address = self.process.allocate_memory(size)
+        address = self.process.allocate_memory(size, preferred_start=preferred_start)
         allocation = Allocation(address, self.process, size)
         self.allocations.append(allocation)
         return allocation

memobj/hook.py

@@ -113,7 +113,7 @@ def get_code(
         """Called when the hook is activated to get the code to put in the hook"""
         raise NotImplementedError()
 
-    def allocate_variable(self, name: str, size: int) -> Allocation:
+    def allocate_variable(self, name: str, size: int, *, preferred_start: int | None = None) -> Allocation:
         """
         Allocate a variable of the specified size for use in the hook, retrievable with get_variable.
 
@@ -122,6 +122,7 @@ def allocate_variable(self, name: str, size: int) -> Allocation:
                 The name of the variable to allocate.
             size: int
                 The size of the memory block to allocate.
+            preferred_start: The preferred start address of the allocation
 
         Returns:
             Allocation
@@ -134,7 +135,7 @@ def allocate_variable(self, name: str, size: int) -> Allocation:
         if self._variables.get("name") is not None:
             raise ValueError(f"Variable {name} is already allocated")
 
-        allocation = self.allocator.allocate(size)
+        allocation = self.allocator.allocate(size, preferred_start=preferred_start)
         self._variables[name] = allocation
         return allocation
 
@@ -338,7 +339,7 @@ def get_hook_tail(self, jump_address: int) -> tuple[list[Instruction], int]:
         else:
             bitness = 32
 
-        decoder = Decoder(bitness, search_bytes, ip=jump_address)
+        decoder = Decoder(bitness, search_bytes, ip=0)
 
         for instruction in decoder:
             # NOTE: this is not a None check, Instruction has special bool() handling
@@ -503,7 +504,7 @@ def get_code(self) -> list[Instruction]:
 @dataclass
 class RegisterCaptureSettings:
     register: RegisterType
-    derefference: bool = True
+    derefference: bool = False
     offset: int = 0
 
 

memobj/process/base.py

@@ -49,13 +49,20 @@ def executable_path(self) -> Path:
         """
         raise NotImplementedError()
 
-    @functools.cached_property
+    @property
     def pointer_format_string(self) -> str:
         if self.process_64_bit:
             return "Q"
         else:
             return "I"
 
+    @property
+    def pointer_type(self) -> Type:
+        if self.process_64_bit:
+            return Type.unsigned8
+        else:
+            return Type.unsigned4
+
     @functools.cached_property
     def pointer_size(self) -> int:
         if self.process_64_bit:
@@ -92,12 +99,13 @@ def from_id(cls, pid: int) -> Self:
     # TODO
     # def itererate_modules(self): ...
 
-    def allocate_memory(self, size: int) -> int:
+    def allocate_memory(self, size: int, *, preferred_start: int | None = None) -> int:
         """
         Allocate <size> amount of memory in the process
 
         Args:
             size: The amount of memory to allocate
+            preferred_start: The preferred start address of the allocation
 
         Returns:
         The start address of the allocated memory

tests/conftest.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from pathlib import Path
 
 import pytest
 
@@ -16,3 +17,21 @@ def process() -> Process:
             return WindowsProcess.from_id(os.getpid())
         case _:
             raise RuntimeError("Unsupported platform")
+
+
+@pytest.fixture(scope="session")
+def test_binaries() -> tuple[Path, Path]:
+    library_root = Path(__file__).parent.parent
+
+    assert library_root.name == "memobj"
+    assert (library_root / "README.md").exists() is True
+
+    release_dir = library_root / "target" / "release"
+
+    exe_path = (release_dir / "inject_target.exe").resolve()
+    dll_path = (release_dir / "test_inject.dll").resolve()
+
+    if not exe_path.exists() or not dll_path.exists():
+        pytest.skip("Test binaries not found")
+
+    return exe_path, dll_path

tests/manual/test_binaries/inject_target/src/main.rs

@@ -1,16 +1,11 @@
 use std::str::FromStr;
 
 fn main() {
-    std::hint::black_box(add_five);
-    std::hint::black_box(create_player);
-    std::hint::black_box(uses_player);
-
-    let value = std::hint::black_box(add_five(60 * 1_000_000));
-
-    // This process just waits, so it can be injected into.
-    //std::thread::sleep(std::time::Duration::from_secs(60 * 5));
-    //std::thread::sleep(std::time::Duration::from_secs(60 * 1_000_000));
-    std::thread::sleep(std::time::Duration::from_secs(value as u64));
+    loop {
+        std::thread::sleep(std::time::Duration::from_secs(1));
+        std::hint::black_box(add_five(1));
+        std::hint::black_box(create_player());
+    }
 }
 
 // static gives us a global address rather than in-lining the value
@@ -52,18 +47,21 @@ impl Player {
 
 // TODO: keep these from getting compiled away
 #[no_mangle]
+#[inline(never)]
 extern "C" fn add_five(value: u32) -> u32 {
     value + FIVE
 }
 
 #[no_mangle]
+#[inline(never)]
 extern "C" fn create_player() -> () {
     let player = Player::new();
 
     uses_player(&player);
 }
 
 #[no_mangle]
+#[inline(never)]
 extern "C" fn uses_player(player: &Player) -> () {
     // cmp with object offset
     if player.health > 100.0 {

tests/manual/test_binaries/test_inject/src/lib.rs

@@ -6,9 +6,6 @@ use windows::Win32::System::SystemServices::DLL_PROCESS_ATTACH;
 use windows::Win32::System::Console::AllocConsole;
 
 
-// #[no_mangle]
-// extern "system" fn DllMain(_: *const u8, _: u32, _: *const u8) -> u32 { 1 }
-
 #[no_mangle]
 extern "system" fn DllMain(
     _hinst_dll: *const u8,
@@ -33,11 +30,14 @@ pub extern "C" fn create_file_at_path(path: *const c_char) -> bool {
 
     let c_str = unsafe { CStr::from_ptr(path) };
 
-    if let Ok(path_str) = c_str.to_str() {
-        if let Ok(mut file) = File::create(path_str) {
-            return file.write_all(b"Injected").is_ok()
-        }
-    }
+    let Ok(path_str) = c_str.to_str() else {
+        return false;
+    };
+
+    let Ok(mut file) = File::create(path_str) else {
+        return false;
+    };
+
+    file.write_all(b"Injected").is_ok()
 
-    false
 }

tests/manual/test_dll_injection.py

@@ -1,37 +1,49 @@
 import subprocess
-import pytest
-from pathlib import Path
+
+from iced_x86 import Register
+import regex
 
 import memobj
+from memobj.utils import ValueWaiter
+from memobj.hook import create_capture_hook, RegisterCaptureSettings
 
-import pytest
 
+def test_dll_injection(test_binaries):
+    exe_path, dll_path = test_binaries
 
-def test_dll_injection():
-    # there is probably a better way to get this
-    library_root = Path(__file__).parent.parent.parent
+    proc = subprocess.Popen([str(exe_path)])
+    try:
+        process = memobj.WindowsProcess.from_id(proc.pid)
+        assert process.inject_dll(dll_path), "DLL injection failed"
+
+        module = process.get_module_named(dll_path.name)
+        assert module is not None, "Failed to find injected DLL in remote process"
+    finally:
+        proc.terminate()
 
-    assert library_root.name == "memobj"
-    assert (library_root / "README.md").exists() is True
 
-    dll_path = (
-        library_root / "target/release/test_inject.dll"
-    ).resolve()
-    exe_path = (
-        library_root / "target/release/inject_target.exe"
-    ).resolve()
+def test_create_capture_hook(test_binaries):
+    exe_path, _ = test_binaries
 
-    if not dll_path.exists():
-        pytest.skip(f"Test DLL not found at {dll_path}")
-    if not exe_path.exists():
-        pytest.skip(f"Test EXE not found at {exe_path}")
+    # TODO: allow passing addresses so we can do a symbol lookup
+    PlayerCaptureHook = create_capture_hook(
+        pattern=regex.escape(bytes.fromhex("48 83 EC 28 F3 0F 10 41 04 0F 2E 05 40 52 01 00 76 0D 8B 09 E8 17 FF FF FF 01 05 21 D0 01 00 48 83 C4 28 C3")),
+        module="inject_target.exe",
+        bitness=64,
+        register_captures=[RegisterCaptureSettings(Register.RCX, derefference=False)],
+    )
 
     proc = subprocess.Popen([str(exe_path)])
     try:
         process = memobj.WindowsProcess.from_id(proc.pid)
-        assert process.inject_dll(dll_path), "DLL injection failed"
+        hook = PlayerCaptureHook(process)
+        hook.activate()
+        rcx_capture = hook.get_variable("RCX_capture")
 
-        module = process.get_module_named(dll_path.name)
-        assert module is not None, "Failed to find injected DLL in remote process"
+        waiter = ValueWaiter(lambda: rcx_capture.read_typed(process.pointer_type))
+        address = waiter.wait_for_value(0, inverse=True, timeout=60)
+
+        assert address != 0
     finally:
         proc.terminate()
+