Skip to content

Instantly share code, notes, and snippets.

@blark
Last active April 2, 2020 02:53
Show Gist options
  • Star 2 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save blark/46d6befc51909c7db4e0 to your computer and use it in GitHub Desktop.
Save blark/46d6befc51909c7db4e0 to your computer and use it in GitHub Desktop.
A really basic completer for prompt_toolkit
from __future__ import print_function, unicode_literals
from prompt_toolkit.completion import Completer, Completion
class PwnCompleter(Completer):
cmds = {'quit': None,
'exit': None,
'use': ['/foo/bar', '/herp/derp', '/baz/qux'],
'info': ['test', '123'],
'show': None,
'shell': None}
meta_dict = {'exit': 'Quit the program',
'shell': 'Drop to a shell' }
def __init__(self, modules):
self.modules = modules
self.categories = list(modules.keys())
self.base_cmds = list(self.cmds.keys())
def get_completions(self, document, complete_event):
word_before_cursor = document.get_word_before_cursor(WORD=True).lower()
all_text = document.text_before_cursor
try:
last_word = all_text.split()[-1]
first_word = all_text.split()[0]
except IndexError:
first_word = None
# if there is no text assign base completions
completions = self.base_cmds
if first_word:
try:
if self.cmds[first_word] is not None:
completions = self.cmds[first_word]
else:
completions = []
except KeyError:
completions = self.base_cmds
completions.sort()
for item in completions:
if word_before_cursor in item:
display_meta = self.meta_dict.get(item, '')
yield Completion(item,
-len(word_before_cursor),
display_meta=display_meta)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment