From ecd725bc724609adc028b97caa1ebae24e9971b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Gaczkowski?= Date: Mon, 30 Mar 2026 14:12:42 +0100 Subject: [PATCH] First draft of new autocomplete Replace old dodgy readline autocomplete with prompt_toolkit --- pyproject.toml | 3 +- requirements.txt | 3 +- src/drozer/console/session.py | 52 +++- src/drozer/modules/common/path_completion.py | 2 +- src/reversec/common/cmd_ext.py | 310 +++++++++++++------ src/reversec/common/path_completion.py | 4 +- 6 files changed, 258 insertions(+), 116 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f3aa592..fe951ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,8 @@ dependencies = ["protobuf>=4.25.2", "service-identity", "distro", "pyyaml", - "requests"] + "requests", + "prompt_toolkit>=3.0.0"] requires-python = ">=3.8" authors = [ { name = "reversec", email = "drozer@reversec.com" }, diff --git a/requirements.txt b/requirements.txt index 6df77e5..e112646 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ twisted>=18.9.0 service-identity distro pyyaml -requests \ No newline at end of file +requests +prompt_toolkit>=3.0.0 \ No newline at end of file diff --git a/src/drozer/console/session.py b/src/drozer/console/session.py index f45ea53..2b31915 100644 --- a/src/drozer/console/session.py +++ b/src/drozer/console/session.py @@ -77,22 +77,48 @@ class Session(cmd.Cmd): def completemodules(self, text): """ - Provides readline auto-completion for drozer module names. + Provides tab-completion for drozer module names. + + Returns one namespace level at a time so the user can drill down + step by step (e.g. "" -> ["app.", "tools."] -> ["app.package."] -> + ["app.package.info"]). Namespace suggestions end with a trailing dot + to signal to the completer that there are more levels below; leaf + module names have no trailing dot. """ - modules = self.modules.all(permissions=self.permissions(), prefix=self.__base) - - if self.__base == "": - modules = list(filter(lambda m: m.startswith(text), modules)) - elif text.startswith("."): - modules = list(filter(lambda m: m.startswith(text[1:]), modules)) - else: - modules = map(lambda m: m[len(self.__base):], list(filter(lambda m: m.startswith(self.__base + text), modules))) - - #if len(modules) == 1 and text == modules[0]: - # return [] + all_modules = self.modules.all(permissions=self.permissions(), prefix=self.__base) - return modules + if self.__base == "": + matching = list(filter(lambda m: m.startswith(text), all_modules)) + elif text.startswith("."): + matching = list(filter(lambda m: m.startswith(text[1:]), all_modules)) + else: + matching = list(map( + lambda m: m[len(self.__base):], + filter(lambda m: m.startswith(self.__base + text), all_modules) + )) + + # How many namespace levels has the user already typed? + # "app.package" has 1 dot → committed to depth 1 → show depth-2 suggestions. + depth = text.count('.') + target_depth = depth + 1 + + seen = set() + result = [] + for m in matching: + parts = m.split('.') + if len(parts) > target_depth: + # More levels exist: suggest only up to the next level, + # with a trailing dot so the completer knows it's a namespace. + suggestion = '.'.join(parts[:target_depth]) + '.' + else: + # This is a leaf module at or within the target depth. + suggestion = m + if suggestion not in seen: + seen.add(suggestion) + result.append(suggestion) + + return result def completenamespaces(self, text): """ diff --git a/src/drozer/modules/common/path_completion.py b/src/drozer/modules/common/path_completion.py index 8a28937..6e65dd4 100644 --- a/src/drozer/modules/common/path_completion.py +++ b/src/drozer/modules/common/path_completion.py @@ -10,5 +10,5 @@ def on_agent(path, context): folder, search_path = get_folder_and_search_path(path, "/") folders = context.listFiles(folder) - return [s.replace(" ", "\ ") for s in get_suggestions(folder, search_path, map(lambda f: str(f), folders), "/", True)] + return [s.replace(" ", "\\ ") for s in get_suggestions(folder, search_path, map(lambda f: str(f), folders), "/", True)] diff --git a/src/reversec/common/cmd_ext.py b/src/reversec/common/cmd_ext.py index 837a7ef..3409ee4 100644 --- a/src/reversec/common/cmd_ext.py +++ b/src/reversec/common/cmd_ext.py @@ -3,19 +3,18 @@ import os from platform import platform from drozer import meta -# readline works on linux/mac -# pyreadline3 works on windows -# or alternatively if you have none of those that's also fine, you just won't get tab-complete -has_readline = False +has_prompt_toolkit = False try: - import readline - has_readline = True -except ModuleNotFoundError: - try: - from pyreadline3 import Readline as readline - has_readline = True - except: - pass + from prompt_toolkit import PromptSession + from prompt_toolkit.completion import Completer, Completion + from prompt_toolkit.history import FileHistory, InMemoryHistory + from prompt_toolkit.auto_suggest import AutoSuggestFromHistory + from prompt_toolkit.formatted_text import HTML + from prompt_toolkit.key_binding import KeyBindings + from prompt_toolkit.filters import has_completions as _has_completions + has_prompt_toolkit = True +except ImportError: + pass import shlex import sys @@ -25,6 +24,122 @@ from reversec.common import system from reversec.common.text import wrap +if has_prompt_toolkit: + def _make_drozer_key_bindings(): + """ + Replaces prompt_toolkit's default Tab cycling behaviour with + "apply and drill down": + + - When the completion menu is open, Tab commits the currently + highlighted completion (or the first one if nothing is highlighted) + and immediately applies it to the buffer. + - For namespace completions (ending in '.'), complete_while_typing + then refreshes the dropdown to the next level automatically. + - For leaf completions the module name is filled in; Enter submits. + - Arrow keys still navigate the menu without committing anything, + so the user can pick a specific item before pressing Tab. + - When the menu is not open, Tab opens it (default behaviour). + + Why not intercept Enter instead? complete_while_typing re-runs the + completer the instant Tab would apply a namespace, which creates a + brand-new complete_state with nothing selected — so any Enter filter + based on completion_is_selected is always False by the time the user + presses Enter. Making Tab do the work sidesteps this entirely. + """ + kb = KeyBindings() + + @kb.add('tab', filter=_has_completions, eager=True) + def _tab_apply_completion(event): + buf = event.current_buffer + cs = buf.complete_state + if cs and cs.completions: + comp = cs.current_completion or cs.completions[0] + buf.apply_completion(comp) + + return kb + + _DROZER_KEY_BINDINGS = _make_drozer_key_bindings() + + class DrozerCompleter(Completer): + """ + Bridges prompt_toolkit's Completer interface to the existing cmd-style + complete_ / completenames dispatch chain. + + All the underlying completion methods (complete_run, complete_cd, + completenames, ArgumentParserCompleter etc.) use the clean + (text, line, begidx, endidx) signature with no readline dependency, + so they plug straight in here. + + WORD=True is critical: it splits only on whitespace, so dotted module + names like "app.package.info" are treated as a single token and passed + intact to the underlying completion functions. + """ + + def __init__(self, cmd_instance): + self._cmd = cmd_instance + + def get_completions(self, document, complete_event): + text_before = document.text_before_cursor + line = text_before.lstrip() + stripped = len(text_before) - len(line) + # WORD=True: only split on whitespace, so "app.package" stays whole + text = document.get_word_before_cursor(WORD=True) + begidx = len(text_before) - len(text) - stripped + endidx = len(text_before) - stripped + + try: + if begidx > 0: + # Output redirection: delegate to completefilename after ">" + if ">" in line and begidx > line.index(">"): + matches = self._cmd.completefilename(text, line, begidx, endidx) or [] + else: + command = self._cmd.parseline(line)[0] + if not command: + compfunc = self._cmd.completedefault + else: + compfunc = getattr(self._cmd, 'complete_' + command, + self._cmd.completedefault) + matches = compfunc(text, line, begidx, endidx) or [] + else: + matches = self._cmd.completenames(text, line, begidx, endidx) or [] + except Exception: + return + + for m in matches: + if m.endswith('.'): + # Namespace suggestion from completemodules: trailing dot signals + # "there are more levels below". Display without dot so the + # dropdown looks clean; insert WITH dot so the next completion + # call immediately shows the next level. + yield Completion(m, start_position=-len(text), display=m.rstrip('.')) + elif m.endswith(os.path.sep): + # Filesystem directory + yield Completion(m, start_position=-len(text)) + else: + # Leaf module, command name, or flag: add trailing space + yield Completion(m.rstrip() + ' ', start_position=-len(text)) + + class _ModuleCompleter(Completer): + """ + Wraps a readline-style (text, state) -> str|None completer callable + for use as a prompt_toolkit Completer. Used for module-provided + completers pushed via push_completer(). + """ + + def __init__(self, readline_fn): + self._fn = readline_fn + + def get_completions(self, document, complete_event): + text = document.get_word_before_cursor() + state = 0 + while True: + result = self._fn(text, state) + if result is None: + break + yield Completion(result, start_position=-len(text)) + state += 1 + + class Cmd(cmd.Cmd): """ An extension to cmd.Cmd to provide some advanced functionality. Including: @@ -55,6 +170,12 @@ class Cmd(cmd.Cmd): self.stderr = sys.stderr self.variables = {} + # prompt_toolkit state + self._pt_session = None + self._pt_completer = None + self._pt_completer_stack = [] + self._pt_history_stack = [] + def cmdloop(self, intro=None): """ Repeatedly issue a prompt, accept input, parse an initial prefix @@ -62,8 +183,14 @@ class Cmd(cmd.Cmd): the remainder of the line as argument. """ self.preloop() - if self.use_rawinput and self.completekey: - self.push_completer(self.complete, self.history_file) + + if has_prompt_toolkit: + history = FileHistory(self.history_file) if self.history_file else InMemoryHistory() + self._pt_completer = DrozerCompleter(self) + self._pt_session = PromptSession(history=history, key_bindings=_DROZER_KEY_BINDINGS) + self._pt_history_stack = [self.history_file] + self._pt_completer_stack = [] + try: stop = None while not stop: @@ -72,9 +199,21 @@ class Cmd(cmd.Cmd): else: if self.use_rawinput: try: - line = input(self.prompt) + if has_prompt_toolkit and self._pt_session: + line = self._pt_session.prompt( + self.prompt, + completer=self._pt_completer, + complete_while_typing=True, + auto_suggest=AutoSuggestFromHistory(), + bottom_toolbar=self._pt_bottom_toolbar, + ) + else: + line = input(self.prompt) except EOFError: line = 'EOF' + except KeyboardInterrupt: + self.stdout.write('\n') + continue else: self.stdout.write(self.prompt) self.stdout.flush() @@ -89,7 +228,7 @@ class Cmd(cmd.Cmd): stop = self.onecmd(line) stop = self.postcmd(stop, line) except ValueError as e: - if e.message == "No closing quotation": + if str(e) == "No closing quotation": self.stderr.write( "Failed to parse your command, because there were unmatched quotation marks.\n") self.stderr.write( @@ -101,58 +240,29 @@ class Cmd(cmd.Cmd): except Exception as e: print("Loop exception") self.handleException(e) - pass - finally: - if self.use_rawinput and self.completekey: - self.pop_completer() - - def complete(self, text, state): + def _pt_bottom_toolbar(self): """ - Return the next possible completion for 'text'. - - If a command has not been entered, then complete against command list. - Otherwise, try to call complete_ to get list of completions. + Returns contextual help for display in the prompt_toolkit bottom toolbar. + Shows the usage line of the command currently being typed. """ - - if state == 0: - if has_readline: - origline = readline.get_line_buffer() - line = origline.lstrip() - stripped = len(origline) - len(line) - begidx = readline.get_begidx() - stripped - endidx = readline.get_endidx() - stripped - - if begidx > 0: - if ">" in line and begidx > line.index(">"): - self.completion_matches = self.completefilename(text, line, begidx, endidx) - return self.completion_matches[0] - - command = self.parseline(line)[0] - if command == '': - compfunc = self.completedefault - else: - try: - compfunc = getattr(self, 'complete_' + command) - except AttributeError: - compfunc = self.completedefault - else: - compfunc = self.completenames - - matches = compfunc(text, line, begidx, endidx) - if len(matches) == 1 and matches[0].endswith(os.path.sep): - self.completion_matches = matches - else: - self.completion_matches = list(map(lambda s: s + " ", matches)) - try: - return self.completion_matches[state] - except IndexError: - return None - except TypeError: - return None + text = self._pt_session.app.current_buffer.text.lstrip() + except AttributeError: + return "" + command = self.parseline(text)[0] if text else "" + if not command: + return "" + try: + raw_doc = (getattr(self, 'do_' + command).__doc__ or "").strip() + usage = next((l.strip() for l in raw_doc.splitlines() if l.strip()), "") + if usage: + safe = usage.replace("<", "<").replace(">", ">") + return HTML(f"{command} — {safe}") + except AttributeError: + pass + return "" - #TODO implement def completefilename(self, text, line, begidx, endidx): """ Placeholder for a filename autocompletion method, that is invoked by @@ -177,11 +287,11 @@ class Cmd(cmd.Cmd): def do_echo(self, arguments): """ usage: echo LINE - + Prints out how a line will be processed at runtime, performing all variable substitutions. - + Example: - + dz> set P=com.example.app dz> echo run app.package.info -a $P run app.package.info com.example.app @@ -202,11 +312,11 @@ class Cmd(cmd.Cmd): def do_set(self, arguments): """ usage: set NAME=VALUE [NAME=VALUE ...] - + Sets one-or-more variables, that can be used to set values in subsequent commands. - + Example: - + dz> set P=com.example.app dz> run app.package.info -a $P """ @@ -219,7 +329,7 @@ class Cmd(cmd.Cmd): def do_unset(self, arguments): """ usage: unset NAME [NAME ...] - + Removes one-or-more values previously stored in variables. """ @@ -306,41 +416,45 @@ class Cmd(cmd.Cmd): self.checkVer() def push_completer(self, completer, history_file=None): - if has_readline: - self.__completer_stack.append(readline.get_completer()) - readline.set_completer(completer) - readline.set_completer_delims(readline.get_completer_delims().replace("/", "")) + """ + Push a new completer (and optionally switch history file) onto the stack. + Used by modules that need to temporarily override tab-completion. + """ + if not has_prompt_toolkit or self._pt_session is None: + return - if len(self.__history_stack) > 0 and self.__history_stack[-1]: - readline.write_history_file(self.__history_stack[-1]) + self._pt_completer_stack.append(self._pt_completer) + self._pt_history_stack.append(history_file) - self.__history_stack.append(history_file) - readline.clear_history() - if history_file is not None and os.path.exists(history_file): - try: - # In macOS, this line causes a `[Errno 1] Operation not permitted` if there is a `~/.drozer_history` - readline.read_history_file(history_file) - except IOError as e: - if "darwin" in platform().lower() and str(e).strip() == "[Errno 1] Operation not permitted": - print - "Could not access the history file..." - else: - raise e + if completer is self.complete: + self._pt_completer = DrozerCompleter(self) + else: + self._pt_completer = _ModuleCompleter(completer) - readline.parse_and_bind(self.completekey + ": complete") + new_history = FileHistory(history_file) if history_file else InMemoryHistory() + self._pt_session = PromptSession(history=new_history, key_bindings=_DROZER_KEY_BINDINGS) def pop_completer(self): - if has_readline: - if self.__history_stack[-1] != None: - readline.write_history_file(self.__history_stack.pop()) - else: - self.__history_stack.pop() + """ + Restore the previous completer and history from the stack. + """ + if not has_prompt_toolkit or not self._pt_completer_stack: + return - readline.clear_history() - if len(self.__history_stack) > 0 and self.__history_stack[-1]: - readline.read_history_file(self.__history_stack[-1]) + self._pt_completer = self._pt_completer_stack.pop() + self._pt_history_stack.pop() - readline.set_completer(self.__completer_stack.pop()) + prev_history_file = self._pt_history_stack[-1] if self._pt_history_stack else None + history = FileHistory(prev_history_file) if prev_history_file else InMemoryHistory() + self._pt_session = PromptSession(history=history, key_bindings=_DROZER_KEY_BINDINGS) + + def complete(self, text, state): + """ + Stub kept for API compatibility with push_completer() callers that + pass `self.complete` as the completer function. The actual completion + work is done by DrozerCompleter when prompt_toolkit is active. + """ + return None def __build_tee(self, console, destination): """ diff --git a/src/reversec/common/path_completion.py b/src/reversec/common/path_completion.py index 75d56fb..f2c77a7 100644 --- a/src/reversec/common/path_completion.py +++ b/src/reversec/common/path_completion.py @@ -6,12 +6,12 @@ def complete(path, include_files=True): Provides path completion, against files local to the Console. """ if path == "": - path = "/" + path = os.path.abspath(".") folder, search_path = get_folder_and_search_path(path, os.path.sep) folders = os.listdir(folder) - return [s.replace(" ", "\ ") for s in get_suggestions(folder, search_path, folders, os.path.sep, include_files)] + return list(get_suggestions(folder, search_path, folders, os.path.sep, include_files)) def get_folder_and_search_path(path, sep):