diff options
Diffstat (limited to 'idatui/rpc.py')
| -rw-r--r-- | idatui/rpc.py | 132 |
1 files changed, 119 insertions, 13 deletions
diff --git a/idatui/rpc.py b/idatui/rpc.py index 7741f48..717f4df 100644 --- a/idatui/rpc.py +++ b/idatui/rpc.py @@ -59,7 +59,7 @@ METHODS = { "text": "{text,delay_ms?,settle?} type a literal string into the focused input", "goto/open": "{target,delay_ms?} g-prompt to a name or 0xADDR", "rename": "{name,word?,delay_ms?} rename the token under (or 'word') the cursor", - "comment": "{text,delay_ms?} comment the current line", + "comment": "{text} comment the current line (use \\n for newlines)", "retype": "{proto,word?,delay_ms?} set the prototype/type under the cursor", "follow": "{word?} follow the reference under (or 'word') the cursor", "cursor_on": "{word,line?,occurrence?=1} place the cursor on a token", @@ -104,6 +104,33 @@ def _active_widget(app): _MODALS = ("XrefsScreen", "SymbolPalette", "StructEditor", "ConfirmScreen") +#: ``drive raw`` (and any k=v CLI) hands every param through as a *string*. +#: Handlers that did ``int(...)`` coped; the ones that compared directly blew up +#: with e.g. "'<' not supported between instances of 'int' and 'str'". Coerce the +#: known-numeric names once, centrally, instead of at every call site. +_INT_PARAMS = ("lines", "limit", "max", "n", "index", "line", "col", + "occurrence", "delay_ms", "direction", "addr", "count") +_FLOAT_PARAMS = ("timeout",) + + +def _coerce_params(params: dict[str, Any]) -> dict[str, Any]: + out = dict(params) + for k in _INT_PARAMS: + v = out.get(k) + if isinstance(v, str) and v.strip(): + try: + out[k] = int(v, 0) + except ValueError: + pass + for k in _FLOAT_PARAMS: + v = out.get(k) + if isinstance(v, str) and v.strip(): + try: + out[k] = float(v) + except ValueError: + pass + return out + def _modal_snapshot(app) -> dict[str, Any] | None: """Describe the top modal screen, if any, enough to drive it.""" @@ -148,6 +175,14 @@ def _cursor_info(app, w) -> dict[str, Any]: "scroll_y": round(w.scroll_offset.y)} +def _where(app) -> str: + """Short 'name @ 0xea' for error messages that need to say where we ended up.""" + cur = getattr(app, "_cur", None) + if cur is None: + return "nowhere" + return f"{getattr(cur, 'name', '?')} @ {getattr(cur, 'ea', 0):#x}" + + def _readiness(app) -> dict[str, Any]: """Whether the app is drivable yet, and how far function-loading has got. (Cheap: no network — never call client.health() here.)""" @@ -452,9 +487,15 @@ class RpcServer: return {"id": rid, "error": {"message": f"{type(e).__name__}: {msg}"}} # -- composed helpers (semantic verbs) -------------------------------- # - async def _press(self, keys, pred=None, timeout=20.0): + async def _press(self, keys, pred=None, timeout=20.0, what=""): await self.app._press_keys([str(k) for k in keys]) - await settle(self.app, pred, timeout=timeout) + ok = await settle(self.app, pred, timeout=timeout) + if pred is not None and not ok: + # Never report success for an action that did not happen: the caller + # would go on to edit whatever the *previous* location was. + raise TimeoutError( + f"{what or 'action'} did not complete within {timeout}s " + f"(still at {_where(self.app)}); retry with a larger timeout=") return snapshot(self.app) async def _fill_prompt(self, open_key, input_id, value, delay_ms, clear): @@ -466,7 +507,14 @@ class RpcServer: await settle(app, lambda: app.query_one(f"#{input_id}", Input).display, timeout=10) inp = app.query_one(f"#{input_id}", Input) if not inp.display: - raise RuntimeError(f"{input_id!r} prompt did not open (word under cursor?)") + # Say *why*. The old message always blamed the word under the cursor, + # which sent readers hunting for a cursor problem when the real cause + # was usually a modal eating the opening keystroke. + modal = type(app.screen).__name__ + why = (f"modal {modal!r} has focus and ate the {open_key!r} keystroke" + if modal in _MODALS or modal != "Screen" + else "no renameable token under the cursor") + raise RuntimeError(f"{input_id!r} prompt did not open: {why}") if clear: inp.value = "" await app._press_keys(_text_to_keys(value, delay_ms)) @@ -485,8 +533,32 @@ class RpcServer: want = fn.addr if fn else ea return lambda: app._cur is not None and app._cur.ea == want + #: Verbs that drive the *main* app by injecting keystrokes. If a modal is on + #: top it eats those keys, so they must refuse rather than silently no-op. + _NEEDS_NO_MODAL = { + "goto", "open", "rename", "comment", "retype", "follow", "back", + "toggle_view", "hex", "save", "search", "move", "cursor", "cursor_on", + } + #: Modals the driver is expected to interact with (they have their own verbs). + _DRIVABLE_MODALS = {"XrefsScreen", "SymbolPalette", "StructEditor", + "ProjectPalette", "QuitScreen"} + + def _modal_kind(self) -> str | None: + scr = self.app.screen + name = type(scr).__name__ + return name if name in _MODALS or name in self._DRIVABLE_MODALS else None + async def _dispatch(self, method: str, params: dict[str, Any]) -> Any: app = self.app + params = _coerce_params(params) + if method in self._NEEDS_NO_MODAL: + modal = self._modal_kind() + if modal is not None: + raise RuntimeError( + f"modal {modal!r} is on top and will swallow this verb's " + f"keystrokes; dismiss it first (close) or use its own verb " + f"(select/symbols/xrefs). Note: a binary with no entry " + f"function can land in the symbol palette on startup.") if method in (None, "ping"): module = None try: @@ -498,9 +570,25 @@ class RpcServer: if method == "methods": return METHODS if method == "quit": + # Route through the same teardown a human gets, so a dirty database + # is written instead of dropped. `app.exit()` alone skips the dirty + # check entirely, and the caller (pane stop) then kills the pane -- + # which used to destroy a whole session's annotations. + dirty = list(app._dirty_labels()) + save = params.get("save", True) + if isinstance(save, str): + save = save.lower() not in ("0", "false", "no", "") + + def _go(): + if dirty and save: + app._on_quit_choice("save") # saves, then exits + else: + app._on_quit_choice("discard") + # answer first, then tear down (so this response still gets written) - asyncio.get_running_loop().call_later(0.2, app.exit) - return {"ok": True, "quitting": True} + asyncio.get_running_loop().call_later(0.2, _go) + return {"ok": True, "quitting": True, "saving": bool(dirty and save), + "dirty": dirty} if method in _PROGRAM_METHODS and app.program is None: raise ValueError("not ready: still connecting / loading functions") @@ -642,7 +730,15 @@ class RpcServer: target = str(params.get("target", "")) pred = self._goto_target_pred(target) await self._fill_prompt("g", "goto", target, delay, clear=False) - await settle(app, pred, timeout=timeout) + ok = await settle(app, pred, timeout=timeout) + if pred is not None and not ok: + # A goto that silently "succeeds" without moving is worse than an + # error: on a big database the listing build can outrun the + # default timeout, and every subsequent rename/comment then lands + # on the function the caller *used* to be looking at. + raise TimeoutError( + f"goto {target!r} did not land within {timeout}s " + f"(still at {_where(app)}); retry with a larger timeout=") return snapshot(app) if method == "rename": @@ -650,7 +746,13 @@ class RpcServer: await settle(app, timeout=timeout) return snapshot(app) if method == "comment": - await self._fill_prompt("semicolon", "comment", str(params["text"]), delay, + # Comments can be long; skip the per-char delay so the agent isn't + # blocked for seconds watching the typing animation. Also: the + # prompt is single-line, so literal newlines (0x0a) get swallowed by + # the Input widget. The app's _do_comment converts the two-char + # sequence '\n' into a real newline for IDA, so we escape here. + ctext = str(params["text"]).replace("\n", "\\n") + await self._fill_prompt("semicolon", "comment", ctext, 0, clear=True) await settle(app, timeout=timeout) return snapshot(app) @@ -661,7 +763,8 @@ class RpcServer: if method == "follow": depth = len(app._nav) - return await self._press(["enter"], lambda: len(app._nav) > depth, timeout) + return await self._press(["enter"], lambda: len(app._nav) > depth, + timeout, "follow") if method == "back": return await self._press(["escape"], timeout=timeout) if method == "toggle_view": @@ -685,12 +788,14 @@ class RpcServer: return False return False - return await self._press(["tab"], _toggled, timeout) + return await self._press(["tab"], _toggled, timeout, "toggle_view") if method == "hex": - return await self._press(["backslash"], lambda: app._active == "hex", timeout) + return await self._press(["backslash"], lambda: app._active == "hex", + timeout, "hex") if method == "xrefs": return await self._press( - ["x"], lambda: type(app.screen).__name__ == "XrefsScreen", timeout) + ["x"], lambda: type(app.screen).__name__ == "XrefsScreen", + timeout, "xrefs") if method == "symbols": await app._press_keys(["ctrl+n"]) await settle(app, lambda: type(app.screen).__name__ == "SymbolPalette", timeout=10) @@ -701,7 +806,8 @@ class RpcServer: return snapshot(app) if method == "structs": return await self._press( - ["ctrl+t"], lambda: type(app.screen).__name__ == "StructEditor", timeout) + ["ctrl+t"], lambda: type(app.screen).__name__ == "StructEditor", + timeout, "structs") if method == "close": return await self._press(["escape"], timeout=timeout) if method == "save": |