Add autostart feature with per-IDB preference

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()
+	return PltPatcher()