diff --git a/plt_patcher.py b/plt_patcher.py index 7e22d33..7ccf48b 100644 --- a/plt_patcher.py +++ b/plt_patcher.py @@ -3,10 +3,133 @@ # Copyright (c) 2025 GAMMACASE # https://github.com/GAMMACASE/PltPatcher +import json +import os + import idc import idaapi import idautils +import ida_auto import ida_segment +import ida_kernwin + +try: + import ida_diskio +except ImportError: + ida_diskio = None + +CONFIG_FILENAME = "plt_patcher.json" +DEFAULT_AUTOSTART = True + +def _get_user_idadir(): + for module in (idaapi, globals().get("ida_diskio")): + if module is None: + continue + getter = getattr(module, "get_user_idadir", None) + if getter is None: + continue + try: + user_dir = getter() + except Exception as exc: + print(f"[PltPatcher] Failed to query IDA user directory: {exc}") + continue + if user_dir: + return user_dir + + user_dir = os.environ.get("IDAUSR") + if user_dir: + return user_dir + + return os.path.join(os.path.expanduser("~"), ".idapro") + +def _get_config_path(): + return os.path.join(_get_user_idadir(), CONFIG_FILENAME) + +def _get_autostart(): + """Read the autostart preference from the current-user config.""" + path = _get_config_path() + try: + with open(path, "r", encoding="utf-8") as config_file: + config = json.load(config_file) + except FileNotFoundError: + return DEFAULT_AUTOSTART + except Exception as exc: + print(f"[PltPatcher] Failed to read config '{path}': {exc}") + return DEFAULT_AUTOSTART + + autostart = config.get("autostart") + if isinstance(autostart, bool): + return autostart + + print(f"[PltPatcher] Invalid autostart config in '{path}', using default") + return DEFAULT_AUTOSTART + +def _set_autostart(enabled): + """Persist the autostart preference into the current-user config.""" + path = _get_config_path() + try: + config_dir = os.path.dirname(path) + if config_dir: + os.makedirs(config_dir, exist_ok=True) + with open(path, "w", encoding="utf-8") as config_file: + json.dump({"autostart": bool(enabled)}, config_file, indent=2) + config_file.write("\n") + except Exception as exc: + print(f"[PltPatcher] Failed to persist autostart config '{path}': {exc}") + return False + + return True + +CONFIG_ACTION_ID = "plt_patcher:configure" +CONFIG_ACTION_LABEL = "PLT Patcher Configuration" + +class PltPatcherConfigHandler(idaapi.action_handler_t): + def __init__(self, plugin): + idaapi.action_handler_t.__init__(self) + self.plugin = plugin + + def activate(self, ctx): + current = self.plugin.autostart + result = ida_kernwin.ask_yn( + 1 if current else 0, + "Auto-patch PLT section when this user opens ELF64 databases?" + ) + if result == -1: # cancelled + return 0 + new_autostart = result == 1 + if new_autostart != current: + self.plugin.autostart = new_autostart + saved = _set_autostart(new_autostart) + print(f"[PltPatcher] Autostart {'enabled' if new_autostart else 'disabled'}") + if not saved: + print("[PltPatcher] Autostart change applies to this session only") + return 1 + + def update(self, ctx): + return idaapi.AST_ENABLE_ALWAYS + +class PltPatcherUIHooks(ida_kernwin.UI_Hooks): + """Defers menu attachment and autostart until the UI is fully ready.""" + + def __init__(self, plugin): + super().__init__() + self.plugin = plugin + + def ready_to_run(self): + ida_kernwin.attach_action_to_menu( + "Edit/Plugins/", CONFIG_ACTION_ID, idaapi.SETMENU_APP + ) + if self.plugin.autostart: + print("[PltPatcher] Auto-patching PLT section...") + self.plugin.run(0) + self.unhook() + +def _wait_for_auto_analysis(): + if ida_auto.auto_is_ok(): + return + + print("[PltPatcher] Waiting for IDA auto-analysis to finish...") + ida_auto.auto_wait() def get_dynamic_struct(): if get_dynamic_struct.dyn_data is None: @@ -112,6 +235,10 @@ def patch_plt(): idc.set_name(got_plt_offs, f'{func_name}_ptr', idaapi.SN_FORCE) if target_ea is not None: + # Skip if .got.plt entry already points to the target + if idaapi.get_qword(got_plt_offs) == target_ea: + continue + # Patch .got.plt entry to point to extern function idaapi.put_qword(got_plt_offs, target_ea) idaapi.add_dref(got_plt_offs, target_ea, idaapi.dr_O) @@ -132,7 +259,7 @@ def patch_plt(): print(f'!!! Failed to find/create {got_plt_offs:x} [{func_name}] function in exports') class PltPatcher(idaapi.plugin_t): - flags = idaapi.PLUGIN_UNL + flags = idaapi.PLUGIN_FIX comment = 'Plt Patcher' help = 'Patches plt sections when IDA fails' wanted_name = 'Patch Plt Section' @@ -141,18 +268,39 @@ def init(self): if 'ELF64' not in idaapi.get_file_type_name(): return idaapi.PLUGIN_SKIP + self.autostart = _get_autostart() + + ida_kernwin.register_action( + ida_kernwin.action_desc_t( + CONFIG_ACTION_ID, + CONFIG_ACTION_LABEL, + PltPatcherConfigHandler(self), + ) + ) + + self._ui_hooks = PltPatcherUIHooks(self) + self._ui_hooks.hook() + + if self.autostart: + print('[PltPatcher] Plugin loaded, PLT will be auto-patched') + else: + print('[PltPatcher] Plugin loaded, use Edit -> Plugins -> Patch Plt Section to run manually') + return idaapi.PLUGIN_KEEP def run(self, arg): print('Starting patching plt section...') + _wait_for_auto_analysis() patch_plt() print('Plt patcher finished.') def term(self): - pass + if hasattr(self, '_ui_hooks'): + self._ui_hooks.unhook() + ida_kernwin.unregister_action(CONFIG_ACTION_ID) def PLUGIN_ENTRY(): - return PltPatcher() \ No newline at end of file + return PltPatcher()