class CommandCompleter(Completer):
"""
Completer for command names.
"""
def __init__(self):
# Completer for full command names.
self._command_completer = WordCompleter(
sorted(COMMANDS_TO_HANDLERS.keys()),
ignore_case=True, WORD=True, match_middle=True)
# Completer for aliases.
self._aliases_completer = WordCompleter(
sorted(ALIASES.keys()),
ignore_case=True, WORD=True, match_middle=True)
def get_completions(self, document, complete_event):
# First, complete on full command names.
found = False
for c in self._command_completer.get_completions(document, complete_event):
found = True
yield c
# When no matches are found, complete aliases instead.
# The completion however, inserts the full name.
if not found:
for c in self._aliases_completer.get_completions(document, complete_event):
full_name = ALIASES.get(c.display)
yield Completion(full_name,
start_position=c.start_position,
display='%s (%s)' % (c.display, full_name))
class JolokiaCliCompleter(Completer):
def __init__(self, mbeans):
self.mbean_dict = {mbean.absolute_object_name: mbean for mbean in mbeans}
self.mbeans_with_attributes_completer = WordCompleter(
[mbean.absolute_object_name for mbean in mbeans if len(mbean.attributes) > 0])
def get_completions(self, document, complete_event):
text_before_cursor = document.text_before_cursor
words_before_cursor = text_before_cursor.split(" ")
if len(words_before_cursor) > 1:
# Command is already complete
command = words_before_cursor[0]
if command in ["read"]:
if len(words_before_cursor) == 2:
for completion in self.mbeans_with_attributes_completer.get_completions(document, complete_event):
yield completion
elif len(words_before_cursor) == 3:
absolute_object_name = words_before_cursor[1]
if absolute_object_name in self.mbean_dict:
attributes = self.mbean_dict[absolute_object_name].attributes
attribute_completer = WordCompleter([attribute.attribute_name for attribute in attributes])
for completion in attribute_completer.get_completions(document, complete_event):
yield completion
elif len(words_before_cursor) == 1:
for possible_command in COMMANDS:
if possible_command.startswith(words_before_cursor[0]):
yield Completion(possible_command, -len(possible_command))
class ExecutableCompleter(Completer):
"""
Complete only excutable files in the current path.
"""
def __init__(self):
self.pathcompleter = PathCompleter(
only_directories=False,
expanduser=True)
self.wordcompleter = WordCompleter(CMDEXE_INT_CMDS + _InternalCmds.keys(), ignore_case=True)
def get_completions(self, document, complete_event):
text_prefix = document.text_before_cursor
# windows cmd.exe command
for completion in self.wordcompleter.get_completions(
document, complete_event):
yield completion
# executeable in PATH
for _dir in os.environ["PATH"].split(os.pathsep):
if not os.path.exists(_dir):
continue
for f in os.listdir(_dir):
if f.lower().startswith(text_prefix.lower()) \
and os.path.isfile(os.path.join(_dir, f)) \
and os.path.splitext(f)[1].lower() in EXE_EXTS:
yield Completion(f, -len(text_prefix), display=f)
# current dir files
for completion in self.pathcompleter.get_completions(
document, complete_event):
yield completion
def get_completions_for_parts(parts, last_part, complete_event, pymux):
completer = None
# Resolve aliases.
if len(parts) > 0:
parts = [ALIASES.get(parts[0], parts[0])] + parts[1:]
if len(parts) == 0:
# New command.
completer = _command_completer
elif len(parts) >= 1 and last_part.startswith('-'):
flags = get_option_flags_for_command(parts[0])
completer = WordCompleter(sorted(flags), WORD=True)
elif len(parts) == 1 and parts[0] in ('set-option', 'set-window-option'):
options = pymux.options if parts[0] == 'set-option' else pymux.window_options
completer = WordCompleter(sorted(options.keys()), sentence=True)
elif len(parts) == 2 and parts[0] in ('set-option', 'set-window-option'):
options = pymux.options if parts[0] == 'set-option' else pymux.window_options
option = options.get(parts[1])
if option:
completer = WordCompleter(sorted(option.get_all_values(pymux)), sentence=True)
elif len(parts) == 1 and parts[0] == 'select-layout':
completer = _layout_type_completer
elif len(parts) == 1 and parts[0] == 'send-keys':
completer = _keys_completer
elif parts[0] == 'bind-key':
if len(parts) == 1:
completer = _keys_completer
elif len(parts) == 2:
completer = _command_completer
# Recursive, for bind-key options.
if parts and parts[0] == 'bind-key' and len(parts) > 2:
for c in get_completions_for_parts(parts[2:], last_part, complete_event, pymux):
yield c
if completer:
for c in completer.get_completions(Document(last_part), complete_event):
yield c
class AutoCommand(Command):
def __init__(self, fn):
self._built_in = False
self._fn = fn
if not callable(fn):
raise ValueError("fn argument must be a callable")
self._obj_metadata = inspect_object(fn)
self._is_super_command = len(self.metadata.subcommands) > 0
self._subcommand_names = []
# We never expect a function to be passed here that has a self argument
# In that case, we should get a bound method
if "self" in self.metadata.arguments and not inspect.ismethod(self._fn):
raise ValueError(
"Expecting either a function (eg. bar) or "
"a bound method (eg. Foo().bar). "
"You passed what appears to be an unbound method "
"(eg. Foo.bar) it has a 'self' argument: %s"
% function_to_str(fn)
)
if not self.metadata.command:
raise ValueError(
"function or class {} needs to be annotated with "
"@command".format(function_to_str(fn))
)
# If this is a super command, we need a completer for sub-commands
if self.super_command:
self._commands_completer = WordCompleter(
[], ignore_case=True, sentence=True
)
for _, inspection in self.metadata.subcommands:
_sub_name = inspection.command.name
self._commands_completer.words.append(_sub_name)
self._commands_completer.meta_dict[_sub_name] = dedent(
inspection.command.help
).strip()
self._subcommand_names.append(_sub_name)
@property
def metadata(self) -> FunctionInspection:
"""
The Inspection object of this command. This object contains all the
information required by AutoCommand to understand the command arguments
type information, help messages, aliases, and attributes.
"""
return self._obj_metadata
def _create_subcommand_obj(self, key_values):
"""
Instantiates an object of the super command class, passes the right
arguments and returns a dict with the remaining unused arguments
"""
kwargs = {
k: v
for k, v in get_arguments_for_inspection(
self.metadata, key_values
).items()
if v is not None
}
remaining = {
k: v for k, v in key_values.items() if k not in kwargs.keys()
}
return self._fn(**kwargs), remaining
def run_interactive(self, cmd, args, raw):
try:
args_metadata = self.metadata.arguments
parsed = parser.parse(args, expect_subcommand=self.super_command)
# prepare args dict
parsed_dict = parsed.asDict()
args_dict = parsed.kv.asDict()
key_values = parsed.kv.asDict()
command_name = cmd
# if this is a super command, we need first to create an instance of
# the class (fn) and pass the right arguments
if self.super_command:
subcommand = parsed_dict.get("__subcommand__")
if not subcommand:
cprint(
"A sub-command must be supplied, valid values: "
"{}".format(", ".join(self._get_subcommands())),
"red",
)
return 2
sub_inspection = self.subcommand_metadata(subcommand)
if not sub_inspection:
cprint(
"Invalid sub-command '{}', valid values: "
"{}".format(
subcommand, ", ".join(self._get_subcommands())
),
"red",
)
return 2
instance, remaining_args = self._create_subcommand_obj(
args_dict
)
assert instance
args_dict = remaining_args
key_values = copy.copy(args_dict)
args_metadata = sub_inspection.arguments
attrname = self._find_subcommand_attr(subcommand)
command_name = subcommand
assert attrname is not None
fn = getattr(instance, attrname)
else:
# not a super-command, use use the function instead
fn = self._fn
positionals = (
parsed_dict["positionals"] if parsed.positionals != "" else []
)
# We only allow positionals for arguments that have positional=True
# ِ We filter out the OrderedDict this way to ensure we don't lose the
# order of the arguments. We also filter out arguments that have
# been passed by name already. The order of the positional arguments
# follows the order of the function definition.
can_be_positional = self._positional_arguments(
args_metadata, args_dict.keys()
)
if len(positionals) > len(can_be_positional):
if len(can_be_positional) == 0:
err = "This command does not support positional arguments"
else:
# We have more positionals than we should
err = (
"This command only supports ({}) positional arguments, "
"namely arguments ({}). You have passed {} arguments ({})"
" instead!"
).format(
len(can_be_positional),
", ".join(can_be_positional.keys()),
len(positionals),
", ".join(positionals),
)
cprint(err, "red")
return 2
# constuct key_value dict from positional arguments.
args_from_positionals = {
key: value for value, key in zip(positionals, can_be_positional)
}
# update the total arguments dict with the positionals
args_dict.update(args_from_positionals)
# Run some validations on number of arguments provided
# do we have keys that are supplied in both positionals and
# key_value style?
duplicate_keys = set(args_from_positionals.keys()).intersection(
set(key_values.keys())
)
if duplicate_keys:
cprint(
"Arguments '{}' have been passed already, cannot have"
" duplicate keys".format(list(duplicate_keys)),
"red",
)
return 2
# check for verbosity override in kwargs
ctx = context.get_context()
old_verbose = ctx.args.verbose
if "verbose" in args_dict:
ctx.set_verbose(args_dict["verbose"])
del args_dict["verbose"]
del key_values["verbose"]
# do we have keys that we know nothing about?
extra_keys = set(args_dict.keys()) - set(args_metadata)
if extra_keys:
cprint(
"Unknown argument(s) {} were"
" passed".format(list(extra_keys)),
"magenta",
)
return 2
# is there any required keys that were not resolved from positionals
# nor key_values?
missing_keys = set(args_metadata) - set(args_dict.keys())
if missing_keys:
required_missing = []
for key in missing_keys:
if not args_metadata[key].default_value_set:
required_missing.append(key)
if required_missing:
cprint(
"Missing required argument(s) {} for command"
" {}".format(required_missing, command_name),
"yellow",
)
return 3
# convert expected types for arguments
for key, value in args_dict.items():
target_type = args_metadata[key].type
if target_type is None:
target_type = str
try:
new_value = apply_typing(value, target_type)
except ValueError as e:
fn_name = function_to_str(target_type, False, False)
cprint(
'Cannot convert value "{}" to {} on argument {}'.format(
value, fn_name, key
),
"yellow",
)
return 4
else:
args_dict[key] = new_value
# Validate that arguments with `choices` are supplied with the
# acceptable values.
for arg, value in args_dict.items():
choices = args_metadata[arg].choices
if choices:
# Validate the choices in the case of values and list of
# values.
if issubclass_(args_metadata[arg].type, typing.List):
bad_inputs = [v for v in value if v not in choices]
if bad_inputs:
cprint(
f"Argument '{arg}' got an unexpected "
f"value(s) '{bad_inputs}'. Expected one "
f"or more of {choices}.",
"red",
)
return 4
elif value not in choices:
cprint(
f"Argument '{arg}' got an unexpected value "
f"'{value}'. Expected one of "
f"{choices}.",
"red",
)
return 4
# arguments appear to be fine, time to run the function
try:
# convert argument names back to match the function signature
args_dict = {
args_metadata[k].arg: v for k, v in args_dict.items()
}
if inspect.iscoroutinefunction(fn):
loop = asyncio.get_event_loop()
ret = loop.run_until_complete(fn(**args_dict))
else:
ret = fn(**args_dict)
ctx.set_verbose(old_verbose)
except Exception as e:
cprint("Error running command: {}".format(str(e)), "red")
cprint("-" * 60, "yellow")
traceback.print_exc(file=sys.stderr)
cprint("-" * 60, "yellow")
return 1
return ret
except CommandParseError as e:
cprint("Error parsing command", "red")
cprint(cmd + " " + args, "white", attrs=["bold"])
cprint((" " * (e.col + len(cmd))) + "^", "white", attrs=["bold"])
cprint(str(e), "yellow")
return 1
def _positional_arguments(self, args_metadata, filter_out):
positionals = OrderedDict()
for k, v in args_metadata.items():
if v.positional and k not in filter_out:
positionals[k] = v
return positionals
def subcommand_metadata(self, name: str) -> FunctionInspection:
assert self.super_command
subcommands = self.metadata.subcommands
for _, inspection in subcommands:
if inspection.command.name == name:
return inspection
def _find_subcommand_attr(self, name):
assert self.super_command
subcommands = self.metadata.subcommands
for attr, inspection in subcommands:
if inspection.command.name == name:
return attr
# be explicit about returning None for readability
return None
def _get_subcommands(self) -> Iterable[str]:
assert self.super_command
return [
inspection.command.name
for _, inspection in self.metadata.subcommands
]
def _kwargs_for_fn(self, fn, args):
return {
k: v
for k, v in get_arguments_for_command(fn, args).items()
if v is not None
}
def run_cli(self, args):
# if this is a super-command, we need to dispatch the call to the
# correct function
kwargs = self._kwargs_for_fn(self._fn, args)
try:
if self._is_super_command:
# let's instantiate an instance of the klass
instance = self._fn(**kwargs)
# we need to find the actual method we want to call, in addition to
# this we need to extract the correct kwargs for this method
# find which function it is in the sub commands
attrname = self._find_subcommand_attr(args._subcmd)
assert attrname is not None
fn = getattr(instance, attrname)
kwargs = self._kwargs_for_fn(fn, args)
else:
fn = self._fn
if inspect.iscoroutinefunction(fn):
# execute in an event loop
loop = asyncio.get_event_loop()
return loop.run_until_complete(fn(**kwargs))
else:
return fn(**kwargs)
except Exception as e:
cprint("Error running command: {}".format(str(e)), "red")
cprint("-" * 60, "yellow")
traceback.print_exc(file=sys.stderr)
cprint("-" * 60, "yellow")
return 1
@property
def super_command(self):
return self._is_super_command
def has_subcommand(self, subcommand):
assert self.super_command
return subcommand.lower() in self._subcommand_names
def add_arguments(self, parser):
register_command(parser, self.metadata)
def get_command_names(self):
command = self.metadata.command
return [command.name] + command.aliases
def get_completions(
self, _: str, document: Document, complete_event: CompleteEvent
) -> Iterable[Completion]:
if self._is_super_command:
exploded = document.text.lstrip().split(" ", 1)
# Are we at the first word? we expect a sub-command here
if len(exploded) <= 1:
return self._commands_completer.get_completions(
document, complete_event
)
state_machine = AutoCommandCompletion(self, document, complete_event)
return state_machine.get_completions()
def get_help(self, cmd, *args):
help = self.metadata.command.help
return dedent(help).strip() if help else None
class ArsenalCompleter(Completer): # pylint: disable-all
"""
A completer specific to the Arsenal API.
"""
_api_methods = []
_built_ins = []
target_names = []
group_names = []
role_names = []
agent_names = []
user_names = []
def __init__(self, methods, autocomplete):
"""
Constructor for the completer, used to gather API information.
"""
self._api_methods = methods
if "*" in self._api_methods:
self._api_methods = list(filter(lambda x: not x.startswith("_"), dir(CLI)))
self._built_ins = [
"help",
"interact",
"exit",
"reset",
"pyexec",
"py3exec",
"pyscript",
"py3script",
"pyexecGroup",
"py3execGroup",
"py3scriptGroup",
"pyscriptGroup",
]
self.target_names = autocomplete.get("target_names", [])
self.group_names = autocomplete.get("group_names", [])
self.role_names = autocomplete.get("role_names", [])
self.agent_names = autocomplete.get("agent_names", [])
self.user_names = autocomplete.get("user_names", [])
self.auto_completers = {
"help": [WordCompleter(self._api_methods)],
"interact": [WordCompleter(self.target_names)],
"reset": [],
"exit": [],
"pyexec": [WordCompleter(self.target_names)],
"py3exec": [WordCompleter(self.target_names)],
"pyscript": [WordCompleter(self.target_names)],
"py3script": [WordCompleter(self.target_names)],
"pyexecGroup": [WordCompleter(self.group_names)],
"py3execGroup": [WordCompleter(self.group_names)],
"pyscriptGroup": [WordCompleter(self.group_names)],
"py3scriptGroup": [WordCompleter(self.group_names)],
"GetTarget": [WordCompleter(self.target_names)],
"GetGroup": [WordCompleter(self.group_names)],
"CreateAction": [WordCompleter(self.target_names)],
"CreateGroupAction": [WordCompleter(self.group_names)],
"AddGroupRule": [WordCompleter(self.group_names)],
"RemoveGroupRule": [WordCompleter(self.group_names)],
"AddGroupMember": [WordCompleter(self.group_names), WordCompleter(self.target_names)],
"RemoveGroupMember": [
WordCompleter(self.group_names),
WordCompleter(self.target_names),
],
"BlacklistGroupMember": [
WordCompleter(self.group_names),
WordCompleter(self.target_names),
],
"AddRoleMember": [WordCompleter(self.role_names), WordCompleter(self.user_names)],
"RemoveRoleMember": [WordCompleter(self.role_names), WordCompleter(self.user_names)],
"GetUser": [WordCompleter(self.user_names)],
"GetRole": [WordCompleter(self.role_names)],
"RenameTarget": [WordCompleter(self.target_names)],
}
self.api_completer = WordCompleter(list(set(self._api_methods + self._built_ins)), True)
def get_completions(self, document, complete_event):
"""
A function for determining auto-complete results.
"""
words = document.text.split(" ")
if words and (words[0] in self._api_methods or words[0] in self._built_ins):
completers = self.auto_completers.get(words[0])
if completers:
try:
completer = completers[len(words) - 2]
if completer:
yield from (
Completion(
completion.text,
completion.start_position,
display=completion.display,
)
for completion in completer.get_completions(document, complete_event)
)
except IndexError:
pass
else:
yield from (
Completion(completion.text, completion.start_position, display=completion.display)
for completion in self.api_completer.get_completions(document, complete_event)
)
class ArsenalCompleter(Completer): # pylint: disable-all
"""
A completer specific to the Arsenal API.
"""
_api_methods = []
_built_ins = []
target_names = []
group_names = []
role_names = []
agent_names = []
user_names = []
def __init__(self, methods, autocomplete):
"""
Constructor for the completer, used to gather API information.
"""
self._api_methods = methods
if '*' in self._api_methods:
self._api_methods = list(
filter(lambda x: not x.startswith('_'), dir(CLI)))
self._built_ins = ['help', 'interact', 'exit', 'reset']
self.target_names = autocomplete.get('target_names', [])
self.group_names = autocomplete.get('group_names', [])
self.role_names = autocomplete.get('role_names', [])
self.agent_names = autocomplete.get('agent_names', [])
self.user_names = autocomplete.get('user_names', [])
self.auto_completers = {
'help': [WordCompleter(self._api_methods)],
'interact': [WordCompleter(self.target_names)],
'reset': [],
'exit': [],
'GetTarget': [WordCompleter(self.target_names)],
'GetGroup': [WordCompleter(self.group_names)],
'CreateAction': [WordCompleter(self.target_names)],
'CreateGroupAction': [WordCompleter(self.group_names)],
'AddGroupRule': [WordCompleter(self.group_names)],
'RemoveGroupRule': [WordCompleter(self.group_names)],
'AddGroupMember': [
WordCompleter(self.group_names),
WordCompleter(self.target_names),
],
'RemoveGroupMember': [
WordCompleter(self.group_names),
WordCompleter(self.target_names),
],
'BlacklistGroupMember': [
WordCompleter(self.group_names),
WordCompleter(self.target_names)
],
'AddRoleMember': [
WordCompleter(self.role_names),
WordCompleter(self.user_names),
],
'RemoveRoleMember': [
WordCompleter(self.role_names),
WordCompleter(self.user_names),
],
'GetUser': [WordCompleter(self.user_names)],
'GetRole': [WordCompleter(self.role_names)],
'RenameTarget': [WordCompleter(self.target_names)]
}
self.api_completer = WordCompleter(
list(set(self._api_methods + self._built_ins)), True)
def get_completions(self, document, complete_event):
"""
A function for determining auto-complete results.
"""
words = document.text.split(' ')
if words and (words[0] in self._api_methods
or words[0] in self._built_ins):
completers = self.auto_completers.get(words[0])
if completers:
try:
completer = completers[len(words) - 2]
if completer:
yield from (Completion(completion.text,
completion.start_position,
display=completion.display)
for completion in completer.
get_completions(document, complete_event))
except IndexError:
pass
else:
yield from (Completion(completion.text,
completion.start_position,
display=completion.display)
for completion in self.api_completer.get_completions(
document, complete_event))
class CommandsRegistry(object):
"""
A registry that holds all commands implementations and creates a quick
access point for resolving a command string into the corresponding handling
object
"""
def __init__(self, parser, listeners):
self._completer = WordCompleter([], ignore_case=True, sentence=True)
# maps a command to Command Instance
self._cmd_instance_map = {}
# objects interested in receiving messages
self._listeners = []
# argparser so each command can add its options
self._parser = parser
for lst in listeners:
self.register_listener(lst(self))
def _make_human_suggestion(self, suggestions):
if len(suggestions) == 1:
return suggestions[0]
human_string = ", ".join(suggestions[:-1])
return human_string + " or {}".format(suggestions[-1])
def register_command(self, cmd_instance, override=False):
if not isinstance(cmd_instance, Command):
raise TypeError("Invalid command instance, must be an instance of "
"subclass of Command")
cmd_instance.set_command_registry(self)
cmd_keys = cmd_instance.get_command_names()
for cmd in cmd_keys:
if not cmd_instance.get_help(cmd):
cprint(
("[WARNING] The command {} will not be loaded. "
"Please provide a help message by either defining a "
"docstring or filling the help argument in the "
"@command annotation").format(cmd_keys[0]),
"red",
)
return None
cmd_instance.add_arguments(self._parser)
if not override:
conflicts = [
cmd for cmd in cmd_keys if cmd in self._cmd_instance_map
]
if conflicts:
raise ValueError("Some other command instance has registered "
"the name(s) {}".format(conflicts))
if isinstance(cmd_instance, Listener):
self._listeners.append(cmd_instance)
for cmd in cmd_keys:
self._cmd_instance_map[cmd.lower()] = cmd_instance
if cmd not in self._completer.words:
self._completer.words.append(cmd)
self._completer.meta_dict[cmd] = cmd_instance.get_help(cmd)
aliases = cmd_instance.get_cli_aliases()
for alias in aliases:
self._cmd_instance_map[alias.lower()] = cmd_instance
def register_priority_listener(self, instance):
"""
Registers a listener that get the top priority in callbacks
"""
if not isinstance(instance, Listener):
raise TypeError("Only Listeners can be registered")
self._listeners.insert(0, instance)
def register_listener(self, instance):
if not isinstance(instance, Listener):
raise TypeError("Only Listeners can be registered")
self._listeners.append(instance)
def __contains__(self, cmd):
return cmd.lower() in self._cmd_instance_map
def get_completer(self):
return self._completer
def get_all_commands(self):
return set(self._cmd_instance_map.values())
def find_command(self, cmd):
return self._cmd_instance_map.get(cmd.lower())
def find_approx(self, cmd):
"""Finds the closest command to the passed cmd, this is used in case we
cannot find an exact match for the cmd
"""
suggestions = []
for c in self._cmd_instance_map.keys():
dist = Levenshtein.distance(str(cmd), str(c))
if dist <= 2:
suggestions.append(c)
if not suggestions:
return ""
return " Did you mean `{}`?".format(
self._make_human_suggestion(suggestions))
def get_completions(self, document, complete_event):
return self._completer.get_completions(document, complete_event)
def dispatch_message(self, msg, *args, **kwargs):
for mod in self._listeners:
mod.react(msg, *args, **kwargs)
def set_cli_args(self, args):
self._args = args
def get_cli_arg(self, arg):
return getattr(self._args, arg, None)
def get_completions(self, document, complete_event):
print(document)
print(complete_event)
res = WordCompleter.get_completions(self, document, complete_event)
print(res)
return res
