Merge pull request #72 from monkeyman192/gh-71

gh-71: Improve path injection for folders of mods

Author: monkeyman192

pymhf/core/hooking.py

@@ -26,7 +26,7 @@
 from pymhf.core.memutils import _get_binary_info, find_pattern_in_binary, get_addressof
 from pymhf.utils.iced import HAS_ICED, generate_load_stack_pointer_bytes, get_first_jmp_addr
 
-hook_logger = logging.getLogger("HookManager")
+logger = logging.getLogger(__name__)
 
 
 BITS = struct.calcsize("P") * 8
@@ -142,7 +142,7 @@ def _determine_detour_list(self, detour: HookProtocol) -> Optional[list[HookProt
             else:
                 detour_list = self._after_detours
         else:
-            hook_logger.error(f"Detour {detour} has an invalid detour time: {detour._hook_time}")
+            logger.error(f"Detour {detour} has an invalid detour time: {detour._hook_time}")
         return detour_list
 
     def add_detour(self, detour: HookProtocol):
@@ -154,15 +154,15 @@ def add_detour(self, detour: HookProtocol):
 
         if getattr(detour, "_noop", False):
             self._has_noop = True
-            hook_logger.warning(
+            logger.warning(
                 f"The hook {detour._hook_func_name} has been marked as NOOP. If there are multiple detours "
                 "registered for this hook they may behave weirdly and the game may crash if this hook is not "
                 "defined well."
             )
 
         # Determine the detour list to use. If none, then return.
         if (detour_list := self._determine_detour_list(detour)) is None:
-            hook_logger.error(
+            logger.error(
                 f"Unable to assign {detour} to a detour type. Please check it. It has not been added."
             )
             return
@@ -186,11 +186,11 @@ def _one_shot(
                     self._disabled_detours.add(detour)
                     detour_list.remove(self._oneshot_detours[detour])
                 except ValueError:
-                    hook_logger.warning(
+                    logger.warning(
                         f"Had an issue removing one-shot {detour._hook_func_name} from detour list."
                     )
                 except Exception:
-                    hook_logger.error(traceback.format_exc())
+                    logger.error(traceback.format_exc())
 
             self._oneshot_detours[detour] = _one_shot
             detour_list.append(_one_shot)
@@ -250,9 +250,9 @@ def bind(self) -> bool:
             super().__init__(signature=self.signature, target=self.target)
         except cyminhook._cyminhook.Error as e:  # type: ignore
             if e.status == cyminhook._cyminhook.Status.MH_ERROR_ALREADY_CREATED:
-                hook_logger.error("Hook is already created")
-            hook_logger.error(f"Failed to initialize hook {self._name} at 0x{self.target:X}")
-            hook_logger.error(e.status.name[3:].replace("_", " ") + f" ({e})")
+                logger.error("Hook is already created")
+            logger.error(f"Failed to initialize hook {self._name} at 0x{self.target:X}")
+            logger.error(e.status.name[3:].replace("_", " ") + f" ({e})")
             self.state = "failed"
             return False
         self.state = "initialized"
@@ -269,8 +269,8 @@ def _compound_detour(self, *args):
                     ret = r
         except Exception:
             bad_detour = self._before_detours.pop(i)
-            hook_logger.error(f"There was an error with detour {bad_detour}. It has been disabled.")
-            hook_logger.error(traceback.format_exc())
+            logger.error(f"There was an error with detour {bad_detour}. It has been disabled.")
+            logger.error(traceback.format_exc())
             self._disabled_detours.add(bad_detour)
 
         # If we don't have any decorators which NOOP the original function then run as usual.
@@ -301,8 +301,8 @@ def _compound_detour(self, *args):
                 bad_detour = self._after_detours.pop(i)
             else:
                 bad_detour = self._after_detours_with_results.pop(j)
-            hook_logger.error(f"There was an error with detour {bad_detour}. It has been disabled.")
-            hook_logger.error(traceback.format_exc())
+            logger.error(f"There was an error with detour {bad_detour}. It has been disabled.")
+            logger.error(traceback.format_exc())
             self._disabled_detours.add(bad_detour)
 
         if after_ret is not None:
@@ -463,7 +463,7 @@ def NOOP(detour: HookProtocol) -> HookProtocol:
     if getattr(detour, "_hook_time", None) == DetourTime.BEFORE:
         setattr(detour, "_noop", True)
     else:
-        hook_logger.warning(
+        logger.warning(
             "NOOP decorator can only be applied to 'before' hooks."
             "Either change the detour to a before detour, or, if it is, ensure that this decorator is "
             "applied above the hook decorator."
@@ -716,11 +716,11 @@ def register_hook(self, hook: HookProtocol):
                 hook_offset = ctypes.cast(func_ptr, ctypes.c_void_p).value
                 hook_offset_is_absolute = True
             else:
-                hook_logger.error(f"Cannot find {hook_binary} in the import list")
+                logger.error(f"Cannot find {hook_binary} in the import list")
                 return
         elif hook._is_exported_func_hook:
             if _internal.BINARY_PATH is None:
-                hook_logger.error("Current running binary path unknown. Cannot hook exported functions")
+                logger.error("Current running binary path unknown. Cannot hook exported functions")
                 return
             # TODO: This is inefficient. We should only instantiate the "dll" once.
             own_dll = ctypes.WinDLL(_internal.BINARY_PATH)
@@ -731,7 +731,7 @@ def register_hook(self, hook: HookProtocol):
         if hook_offset is not None:
             func_id = FunctionIdentifier(hook_func_name, hook_offset, hook_binary, hook_offset_is_absolute)
         else:
-            hook_logger.error(f"Unable to find offset for {hook_func_name}. Hook will not be registered.")
+            logger.error(f"Unable to find offset for {hook_func_name}. Hook will not be registered.")
             return
 
         if func_id not in self.hooks:
@@ -744,7 +744,7 @@ def register_hook(self, hook: HookProtocol):
                     offset_is_absolute=func_id.is_absolute,
                 )
             except Exception:
-                hook_logger.exception(f"There was an issue creating the func hook for {func_id}")
+                logger.exception(f"There was an issue creating the func hook for {func_id}")
             self._uninitialized_hooks.add(func_id)
             self._hook_id_mapping[hook] = func_id
         self.hooks[func_id].add_detour(hook)
@@ -774,16 +774,16 @@ def initialize_hooks(self) -> int:
                 else:
                     offset = hook.offset
                     prefix = f"{hook._binary}+"
-                hook_logger.info(f"Enabled hook for {hook_func_id.name} at {prefix}0x{offset:X}")
+                logger.debug(f"Enabled hook for {hook_func_id.name} at {prefix}0x{offset:X}")
             except Exception:
-                hook_logger.error(f"Unable to enable {hook_func_id.name} because:")
-                hook_logger.error(traceback.format_exc())
+                logger.error(f"Unable to enable {hook_func_id.name} because:")
+                logger.error(traceback.format_exc())
 
             # If any of the hooked functions want to log where they were called from, we need to overwrite
             # part of the trampoline bytes to capture the RSP register.
             if hook_func_id in self._get_caller_detours:
                 if not HAS_ICED:
-                    hook_logger.error(
+                    logger.error(
                         f"Cannot get calling address of {hook_func_id.name} as `iced_x86` package is not "
                         "installed.\nPlease install and try again."
                     )
@@ -806,7 +806,7 @@ def initialize_hooks(self) -> int:
                         data_at_detour[i] = rsp_load_bytes[i]
                     for j in range(len(orig_bytes)):
                         data_at_detour[i + j + 1] = orig_bytes[j]
-                    hook_logger.info(
+                    logger.info(
                         f"The function {hook_func_id.name} has a modified hook to get the calling address."
                     )
 
@@ -817,15 +817,15 @@ def initialize_hooks(self) -> int:
     def _debug_show_states(self):
         # Return the states of all the registered hooks
         for hook_func_id, hook in self.hooks.items():
-            hook_logger.info(f"Functions registered for {hook_func_id.name}:")
+            logger.info(f"Functions registered for {hook_func_id.name}:")
             if hook._before_detours:
-                hook_logger.info("  Before Detours:")
+                logger.info("  Before Detours:")
                 for func in hook._before_detours:
-                    hook_logger.info(f"    {func}")
+                    logger.info(f"    {func}")
             if hook._after_detours:
-                hook_logger.info("  After Detours:")
+                logger.info("  After Detours:")
                 for func in hook._after_detours:
-                    hook_logger.info(f"    {func}")
+                    logger.info(f"    {func}")
 
 
 class Structure(ctypes.Structure):
@@ -925,16 +925,16 @@ def _call(self, *args, **kwargs) -> Optional[R]:
                 try:
                     val = cfunc(*_args)
                 except ctypes.ArgumentError:
-                    hook_logger.error(
+                    logger.error(
                         f"{self._func.__qualname__!r} has function signature {self._funcdef.arg_types} "
                         f"but was called with {_args}"
                     )
                     raise
                 return val
             else:
-                hook_logger.error(f"Unable to call {self._func.__qualname__!r} - Cannot find function.")
+                logger.error(f"Unable to call {self._func.__qualname__!r} - Cannot find function.")
         except Exception:
-            hook_logger.exception(f"There was an exception calling {self._func.__qualname__!r}")
+            logger.exception(f"There was an exception calling {self._func.__qualname__!r}")
             return None
 
     def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Optional[R]:
@@ -962,7 +962,7 @@ def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                         # If it's not a pointer, then we'll assume it's an int and pass the address...
                         return self._call(ctypes.addressof(self._bound_class), *args, **kwargs)
                 except Exception:
-                    hook_logger.exception(f"Failed to call {self._func.__qualname__} with args {args}")
+                    logger.exception(f"Failed to call {self._func.__qualname__} with args {args}")
             else:
                 raise ValueError("Not bound to anything...")
 

pymhf/core/importing.py

@@ -1,14 +1,14 @@
+import ast
 import importlib
 import importlib.util
 import logging
 import os.path as op
 import string
 import sys
-import traceback
 from types import ModuleType
 from typing import Optional
 
-logger = logging.getLogger("pymfh.core.importing")
+logger = logging.getLogger(__name__)
 
 
 VALID_CHARS = string.ascii_letters + string.digits + "_"
@@ -27,6 +27,53 @@ def _clean_name(name: str) -> str:
     return out
 
 
+def _fully_unpack_ast_attr(obj: ast.Attribute) -> str:
+    name = ""
+    _obj = obj
+    while isinstance(_obj, ast.Attribute):
+        if name:
+            name = f"{_obj.attr}.{name}"
+        else:
+            name = _obj.attr
+        _obj = _obj.value
+    else:
+        if isinstance(_obj, ast.Name):
+            name = f"{_obj.id}.{name}"
+    return name
+
+
+def parse_file_for_mod(data: str) -> bool:
+    """Parse the provided data and determine if there is at least one mod class in it."""
+    tree = ast.parse(data)
+    mod_class_name = None
+    for node in tree.body:
+        # First, determine the name the Mod object is imported as.
+        if isinstance(node, ast.Import):
+            for node_ in node.names:
+                if isinstance(node_, ast.alias):
+                    if node_.name in ("pymhf", "pymhf.core.mod_loader"):
+                        mod_class_name = (node_.asname or node_.name) + ".Mod"
+        if isinstance(node, ast.ImportFrom):
+            if node.module in ("pymhf", "pymhf.core.mod_loader"):
+                for node_ in node.names:
+                    if isinstance(node_, ast.alias):
+                        if node_.name == "Mod":
+                            mod_class_name = node_.asname or node_.name
+        # Now, when we go over the class nodes, check the base classes.
+        if isinstance(node, ast.ClassDef):
+            for base in node.bases:
+                # For a simple name, it's easy - just match it.
+                if isinstance(base, ast.Name):
+                    if base.id == mod_class_name:
+                        return True
+                # If it's an attribute it's a bit trickier...
+                elif isinstance(base, ast.Attribute):
+                    resolved_base = _fully_unpack_ast_attr(base)
+                    if resolved_base == mod_class_name:
+                        return True
+    return False
+
+
 def import_file(fpath: str) -> Optional[ModuleType]:
     try:
         module_name = _clean_name(op.splitext(op.basename(fpath))[0])
@@ -44,5 +91,4 @@ def import_file(fpath: str) -> Optional[ModuleType]:
         else:
             print("failed")
     except Exception:
-        logger.error(f"Error loading {fpath}")
-        logger.error(traceback.format_exc())
+        logger.exception(f"Error loading {fpath}")

pymhf/core/logging.py pymhf/core/log_handling.pypymhf/core/logging.py renamed to pymhf/core/log_handling.py

@@ -20,7 +20,7 @@
 BLACKLIST = type, ModuleType, FunctionType
 
 
-mem_logger = logging.getLogger("MemUtils")
+logger = logging.getLogger(__name__)
 
 MEM_ACCESS_R = 0x100  # Read only.
 MEM_ACCESS_RW = 0x200  # Read and Write access.
@@ -276,7 +276,7 @@ def find_pattern_in_binary(
     binary we ran, but it could be some other dll under the same process.)
     """
     if (_cached_offset := cache.offset_cache.get(pattern, binary)) is not None:
-        mem_logger.debug(
+        logger.debug(
             f"Using cached offset 0x{_cached_offset:X} for pattern {pattern}"
             f" and binary {binary or _internal.EXE_NAME}"
         )
@@ -293,6 +293,6 @@ def find_pattern_in_binary(
         _offset = _offset - module.lpBaseOfDll
     # Cache even if there is no result (so we don't repeatedly look for it when it's not there in case there
     # is an issue.)
-    mem_logger.debug(f"Found {pattern} at 0x{_offset:X} for binary {binary}")
+    logger.debug(f"Found {pattern} at 0x{_offset:X} for binary {binary}")
     cache.offset_cache.set(pattern, _offset, binary, True)
     return _offset