trainerbase 4.5.1__tar.gz → 4.6.0__tar.gz

{trainerbase-4.5.1 → trainerbase-4.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trainerbase
-Version: 4.5.1
+Version: 4.6.0
 Summary: Framework for creating game trainers/mods
 Project-URL: Homepage, https://gitlab.com/python-trainers/
 Project-URL: Repository, https://gitlab.com/python-trainers/trainer-base

{trainerbase-4.5.1 → trainerbase-4.6.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "trainerbase"
-version = "4.5.1"
+version = "4.6.0"
 description = "Framework for creating game trainers/mods"
 license = "MIT"
 authors = [

{trainerbase-4.5.1 → trainerbase-4.6.0}/trainerbase/codeinjection.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import suppress
+from functools import lru_cache
 from pathlib import Path
 from subprocess import DEVNULL
 from subprocess import run as run_executable
@@ -15,6 +16,9 @@ from trainerbase.process import pm
 
 
 FASM_EXECUTABLE_PATH: Final[Path] = config.vendor.folder / "FASM.EXE"
+ORIGINAL_CODE_PLACEHOLDER: Final[str] = "__original_code__"
+BASE_ADDRESS_PLACEHOLDER: Final[str] = "__base_address__"
+MOVE_X64_MACROS: Final[str] = "__move_x64__"
 
 
 type Code = bytes | str
@@ -71,8 +75,8 @@ class CodeInjection(AbstractCodeInjection):
             logger.debug("Nothing to eject")
 
 
-class AllocatingCodeInjection(AbstractCodeInjection):  # pylint: disable=too-many-instance-attributes
-    def __init__(  # pylint: disable=too-many-arguments
+class AllocatingCodeInjection(AbstractCodeInjection):
+    def __init__(
         self,
         address: ConvertibleToAddress,
         code_to_inject: Code,
@@ -86,10 +90,7 @@ class AllocatingCodeInjection(AbstractCodeInjection):  # pylint: disable=too-man
 
         self.is_long_x64_jump_needed = is_long_x64_jump_needed
         self.address = ensure_address(address)
-        self.code_to_inject = compile_asm(code_to_inject)
-
-        if self.is_long_x64_jump_needed:
-            self.code_to_inject = compile_asm("pop rax") + self.code_to_inject + compile_asm("push rax")
+        self.code_to_inject = code_to_inject
 
         self.original_code: bytes = b""
 
@@ -101,12 +102,25 @@ class AllocatingCodeInjection(AbstractCodeInjection):  # pylint: disable=too-man
     @override
     def inject(self):
         self.injection_point = self.address.resolve()
-
         self.new_memory_address: int = allocate(self.new_memory_size)
-
         self.original_code = pm.read_bytes(self.injection_point, self.original_code_length)
 
+        code_to_inject = self.code_to_inject
+
+        if isinstance(code_to_inject, str):
+            code_to_inject = preprocess_asm(
+                code_to_inject,
+                original_code=self.original_code,
+                base_address=pm.base_address,
+            )
+
+        if isinstance(code_to_inject, bytearray):
+            raise TypeError("Code to inject must not be bytearray")
+
+        code_to_inject = compile_asm(code_to_inject)
+
         if self.is_long_x64_jump_needed:
+            code_to_inject = compile_asm("pop rax") + code_to_inject + compile_asm("push rax")
             jump_to_new_mem = compile_asm(get_long_x64_jump_to_new_mem_code(self.new_memory_address))
             jumper_length = len(jump_to_new_mem)
             jump_back_address = self.injection_point + jumper_length - 1
@@ -116,7 +130,7 @@ class AllocatingCodeInjection(AbstractCodeInjection):  # pylint: disable=too-man
             jumper_length = len(jump_to_new_mem)
             jump_back_address = self.injection_point + jumper_length
             jump_back = compile_asm(
-                get_simple_jump_code(jump_back_address, self.new_memory_address + len(self.code_to_inject)),
+                get_simple_jump_code(jump_back_address, self.new_memory_address + len(code_to_inject)),
             )
 
         if jumper_length < self.original_code_length:
@@ -124,7 +138,7 @@ class AllocatingCodeInjection(AbstractCodeInjection):  # pylint: disable=too-man
         elif jumper_length > self.original_code_length:
             raise ValueError(f"Jumper length > original code length: {jumper_length} > {self.original_code_length}")
 
-        code_to_inject = self.code_to_inject + jump_back
+        code_to_inject = code_to_inject + jump_back
 
         logger.debug(f"Injecting {len(code_to_inject)} byte(s) (INJECTION) at 0x{self.new_memory_address:X}")
         pm.write_bytes(self.new_memory_address, code_to_inject, len(code_to_inject))
@@ -161,6 +175,29 @@ class MultipleCodeInjection(AbstractCodeInjection):
             injection.eject()
 
 
+def preprocess_asm(code: str, *, original_code: bytes | None, base_address: int | None) -> str:
+    new_lines = []
+
+    for line in code.splitlines():
+        if original_code is not None:
+            formatted_bytes = ", ".join(map(hex, original_code))
+            formatted_original_code = f"db {formatted_bytes}"
+
+            line = line.replace(ORIGINAL_CODE_PLACEHOLDER, formatted_original_code)
+
+        if base_address is not None:
+            line = line.replace(BASE_ADDRESS_PLACEHOLDER, hex(base_address))
+
+        if MOVE_X64_MACROS in line:
+            _, address, register = line.strip().split()
+            line = get_move_x64_register_to_memory_code(address, register)
+
+        new_lines.append(line)
+
+    return "\n".join(new_lines)
+
+
+@lru_cache
 def compile_asm(code: Code) -> bytes:
     if isinstance(code, str):
         logger.debug(f"Assembling for x{ARCH}:\n{code}")
@@ -204,6 +241,18 @@ def get_long_x64_jump_back_code(address: int) -> str:
     """
 
 
+def get_move_x64_register_to_memory_code(address: int | str, register: str) -> str:
+    temporary_register = "rbx" if register == "rax" else "rax"
+    code = f"""
+        push {temporary_register}
+        mov {temporary_register}, {address}
+        mov [{temporary_register}], {register}
+        pop {temporary_register}
+    """
+
+    return code
+
+
 def safely_eject_all_code_injections() -> None:
     for injection in AbstractCodeInjection.created_code_injections:
         with suppress(MemoryWriteError):

{trainerbase-4.5.1 → trainerbase-4.6.0}/trainerbase/config.py

@@ -14,7 +14,7 @@ DEFAULT_VENDOR_FOLDER_PATHS: Final[list[Path]] = [
 
 
 class Process(BaseModel):
-    names: list[str] = Field(alias=AliasChoices("names", "process_names"))
+    names: list[str] = Field(validation_alias=AliasChoices("names", "process_names"))
     exact_match: bool = True
     ignore_case: bool = False
 
@@ -44,8 +44,8 @@ class Vendor(BaseModel):
 
 
 class Config(BaseModel):
-    process: Process = Field(alias=AliasChoices("process", "pymem"))
-    vendor: Vendor | None = Vendor()
+    process: Process = Field(validation_alias=AliasChoices("process", "pymem"))
+    vendor: Vendor = Vendor()
     logging: dict[str, Any] | None = None  # Must be managed by logging.dictConfig
 
 

{trainerbase-4.5.1 → trainerbase-4.6.0}/trainerbase/memory.py

@@ -207,3 +207,41 @@ def get_module_address(module_name: str, *, delay: float = 2, retry_limit: int |
         sleep(delay)
 
     raise RuntimeError("This code must be unreachable!")
+
+
+def ensure_resolved(
+    address: AbstractAddress,
+    verbose_name: str,
+    *,
+    delay: float = 2,
+    retry_limit: int | None = 4,
+) -> int:
+    """
+    Returns resolved address.
+    Raises RuntimeError if the address was not found and retry_limit is reached.
+
+    Set retry_limit to None or negative value if you need infinite retry cycle.
+    Set retry_limit to 0 if you want to disable retries.
+    """
+
+    logger.info(f"Looking for {verbose_name}...")
+
+    for try_counter in count():
+        try:
+            resolved = address.resolve()
+        except ValueError as e:
+            if try_counter == retry_limit:
+                raise RuntimeError(f"Failed to resolve {verbose_name}, tried {try_counter + 1} time(s).") from e
+
+            logger.info(f"Failed to resolve {verbose_name}. I will retry in {delay} s. Retry limit is {retry_limit}")
+
+            sleep(delay)
+            continue
+
+        logger.info(f"Successfully found {verbose_name} at 0x{resolved:X}")
+        return resolved
+
+    raise RuntimeError("This code must be unreachable!")
+
+
+mainaob = partial(AOBScan, module_name=pm.process_base.name)