Skip to content

Commit 92c00ac

Browse files
authored
Add test app for testing dll injection (#79)
* add test app stuff * remove the formatted single methods * formatting and import sorting * fix missing not in 64 bit check * Merge branch 'dev' of https://github.com/StarrFox/memobj into dev * fix more typing issues * attempt getting module within allocator * switch to using WindowsModule over WindowsModuleInfo in process * add test app * fix unused import * uv install all groups by default * fix tests cache key
1 parent 2b465b6 commit 92c00ac

28 files changed

+742
-344
lines changed

.github/workflows/tests.yml

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ on:
1818

1919
jobs:
2020
test:
21-
2221
runs-on: windows-latest
2322
strategy:
2423
fail-fast: false
@@ -27,13 +26,39 @@ jobs:
2726

2827
steps:
2928
- uses: actions/checkout@v4
29+
30+
- name: Restore cached test binaries
31+
id: restore-cache-test
32+
uses: actions/cache/restore@v4
33+
with:
34+
path: tests/manual/test_inject/target
35+
key: ${{ runner.os }}-${{ hashFiles('tests/manual/test_inject/src/**') }}
36+
37+
- name: Setup rust
38+
if: steps.restore-cache-test.outputs.cache-hit != 'true'
39+
uses: actions-rust-lang/setup-rust-toolchain@v1
40+
41+
- name: Build test binaries
42+
if: steps.restore-cache-test.outputs.cache-hit != 'true'
43+
run: cargo build --release --manifest-path tests/manual/test_inject/Cargo.toml
44+
45+
- name: Cache test binaries
46+
if: steps.restore-cache-test.outputs.cache-hit != 'true'
47+
uses: actions/cache/save@v4
48+
with:
49+
path: tests/manual/test_inject/target
50+
key: ${{ runner.os }}-${{ hashFiles('tests/manual/test_inject/src/**') }}
51+
3052
- name: Set up Python ${{ matrix.python-version }}
3153
uses: actions/setup-python@v5
3254
with:
3355
python-version: ${{ matrix.python-version }}
56+
3457
- name: Install uv
3558
uses: astral-sh/setup-uv@v6
59+
3660
- name: Install dependencies
37-
run: uv sync --all-groups
61+
run: uv sync
62+
3863
- name: Run tests
3964
run: uv run pytest

Cargo.lock

Lines changed: 155 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
[workspace]
2+
members = [
3+
"tests/manual/test_inject"
4+
]

README.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,9 @@ This project uses uv as the build backend and package manager.
4343
- isort . && black .
4444

4545
Optional: A Nix flake provides a dev shell with Python 3.11, just, black, isort, and more:
46+
4647
- nix develop
4748

4849
## Support
4950

50-
discord: https://discord.gg/wcftyYm6qe
51+
discord: <https://discord.gg/wcftyYm6qe>

flake.nix

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
python
3232
just
3333
alejandra
34+
rustc
35+
cargo
3436
python.pkgs.black
3537
python.pkgs.isort
3638
python.pkgs.vulture

justfile

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,3 +46,18 @@ format:
4646
isort .
4747
black .
4848
alejandra .
49+
50+
# build test dll
51+
build-test-dll:
52+
if (!(Test-Path "tests/manual/test_inject/target/release/test_inject.dll")) { cargo build --release --manifest-path tests/manual/test_inject/Cargo.toml }
53+
54+
# build test exe
55+
build-test-exe:
56+
if (!(Test-Path "tests/manual/test_inject/target/release/inject_target.dll")) { cargo build --release --manifest-path tests/manual/test_inject/Cargo.toml --bin inject_target }
57+
58+
# run manual tests
59+
manual-test: build-test-dll build-test-exe
60+
uv run pytest -rs --run-manual tests/manual/
61+
62+
# run all tests
63+
all-tests: test manual-test

memobj/__init__.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1+
import logging
2+
3+
from .allocation import Allocation, Allocator
4+
from .object import MemoryObject
15
from .process import Process, WindowsProcess
26
from .property import *
3-
from .object import MemoryObject
4-
from .allocation import Allocator, Allocation
5-
6-
import logging
77

88
logger = logging.getLogger(__name__)
99
logger.addHandler(logging.NullHandler())

memobj/allocation.py

Lines changed: 43 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,16 @@
1-
from typing import TYPE_CHECKING, Self
1+
from typing import TYPE_CHECKING, Any, Self
22

3+
from memobj.utils import Type, get_type_size
34

45
if TYPE_CHECKING:
56
from memobj.process import Process
67

78

89
class Allocation:
9-
def __init__(self, address: int, process: "Process"):
10+
def __init__(self, address: int, process: "Process", size: int):
1011
self.address = address
1112
self.process = process
13+
self.size = size
1214

1315
self._is_closed: bool = False
1416

@@ -32,6 +34,22 @@ def free(self):
3234
self.process.free_memory(self.address)
3335
self._is_closed = True
3436

37+
def read_typed(self, read_type: Type) -> Any:
38+
if (type_size := get_type_size(read_type)) != self.size:
39+
raise ValueError(
40+
f"{read_type} size ({type_size}) does not match allocation size ({self.size})"
41+
)
42+
43+
return self.process.read_typed(self.address, read_type)
44+
45+
def write_typed(self, write_type: Type, value: Any) -> None:
46+
if (type_size := get_type_size(write_type)) != self.size:
47+
raise ValueError(
48+
f"Write type ({type_size}) does not match allocation size ({self.size})"
49+
)
50+
51+
return self.process.write_typed(self.address, write_type, value)
52+
3553

3654
class Allocator:
3755
"""
@@ -61,12 +79,34 @@ def closed(self) -> bool:
6179
return self._is_closed
6280

6381
def allocate(self, size: int) -> Allocation:
82+
"""
83+
Allocates a block of memory for the process.
84+
85+
Allocates a specified block of memory for the associated process, keeping
86+
track of the allocation for management purposes. The allocated memory
87+
is represented as an `Allocation` object and is appended to the list of
88+
current allocations.
89+
90+
Args:
91+
size (int): The size of the memory block to allocate in bytes.
92+
93+
Returns:
94+
Allocation: An object representing the allocated memory block.
95+
"""
6496
address = self.process.allocate_memory(size)
65-
allocation = Allocation(address, self.process)
97+
allocation = Allocation(address, self.process, size)
6698
self.allocations.append(allocation)
6799
return allocation
68100

69101
def close(self):
102+
"""
103+
Closes the allocator, ensuring all current allocations are properly freed and the allocator
104+
is set to a closed state. This method prevents further use of the allocator by marking it
105+
as closed. If the allocator is already closed, an error will be raised.
106+
107+
Raises:
108+
ValueError: If the allocator is already closed before the invocation of this method.
109+
"""
70110
if self._is_closed:
71111
raise ValueError("Cannot close an already closed allocator")
72112

0 commit comments

Comments
 (0)