CLI: Show choices options when no argument provided

Fixes issue #191.
This commit is contained in:
Mike Auty
2020-02-15 13:23:19 +00:00
parent 6a6cf37b02
commit 3a824b519e
2 changed files with 85 additions and 50 deletions
+7 -50
View File
@@ -10,7 +10,6 @@ User interfaces make use of the framework to:
* run the plugin
* display the results
"""
import argparse
import glob
import inspect
@@ -19,13 +18,13 @@ import logging
import os
import sys
import traceback
from typing import Any, Dict, Type, Union
from typing import Dict, Type, Union
from urllib import parse, request
import volatility.plugins
import volatility.symbols
from volatility import framework
from volatility.cli import text_renderer
from volatility.cli import text_renderer, volargparse
from volatility.framework import automagic, constants, contexts, exceptions, interfaces, plugins, configuration
from volatility.framework.configuration import requirements
@@ -83,7 +82,8 @@ class CommandLine(interfaces.plugins.FileConsumerInterface):
renderers = dict([(x.name.lower(), x) for x in framework.class_subclasses(text_renderer.CLIRenderer)])
parser = argparse.ArgumentParser(prog = 'volatility', description = "An open-source memory forensics framework")
parser = volargparse.HelpfulArgParser(prog = 'volatility',
description = "An open-source memory forensics framework")
parser.add_argument("-c",
"--config",
help = "Load the configuration from a json file",
@@ -212,7 +212,9 @@ class CommandLine(interfaces.plugins.FileConsumerInterface):
self.populate_requirements_argparse(parser, amagic.__class__)
configurables_list[amagic.__class__.__name__] = amagic
subparser = parser.add_subparsers(title = "Plugins", dest = "plugin", action = HelpfulSubparserAction)
subparser = parser.add_subparsers(title = "Plugins",
dest = "plugin",
action = volargparse.HelpfulSubparserAction)
for plugin in sorted(plugin_list):
plugin_parser = subparser.add_parser(plugin, help = plugin_list[plugin].__doc__)
self.populate_requirements_argparse(plugin_parser, plugin_list[plugin])
@@ -481,51 +483,6 @@ class CommandLine(interfaces.plugins.FileConsumerInterface):
**additional)
# We shouldn't really steal a private member from argparse, but otherwise we're just duplicating code
class HelpfulSubparserAction(argparse._SubParsersAction):
"""Class to either select a unique plugin based on a substring, or identify
the alternatives."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# We don't want the action self-check to kick in, so we remove the choices list, the check happens in __call__
self.choices = None
def __call__(self, parser, namespace, values, option_string = None):
parser_name = values[0]
arg_strings = values[1:]
# set the parser name if requested
if self.dest is not argparse.SUPPRESS:
setattr(namespace, self.dest, parser_name)
matched_parsers = [name for name in self._name_parser_map if parser_name in name]
if len(matched_parsers) < 1:
msg = 'invalid choice {} (choose from {})'.format(parser_name, ', '.join(self._name_parser_map))
raise argparse.ArgumentError(self, msg)
if len(matched_parsers) > 1:
msg = 'plugin {} matches multiple plugins ({})'.format(parser_name, ', '.join(matched_parsers))
raise argparse.ArgumentError(self, msg)
parser = self._name_parser_map[matched_parsers[0]]
setattr(namespace, 'plugin', matched_parsers[0])
# parse all the remaining options into the namespace
# store any unrecognized options on the object, so that the top
# level parser can decide what to do with them
# In case this subparser defines new defaults, we parse them
# in a new namespace object and then update the original
# namespace for the relevant parts.
subnamespace, arg_strings = parser.parse_known_args(arg_strings, None)
for key, value in vars(subnamespace).items():
setattr(namespace, key, value)
if arg_strings:
vars(namespace).setdefault(argparse._UNRECOGNIZED_ARGS_ATTR, [])
getattr(namespace, argparse._UNRECOGNIZED_ARGS_ATTR).extend(arg_strings)
def main():
"""A convenience function for constructing and running the
:class:`CommandLine`'s run method."""
+78
View File
@@ -0,0 +1,78 @@
import argparse
import gettext
import re
# This effectively overrides/monkeypatches the core argparse module to provide more helpful output around choices
# We shouldn't really steal a private member from argparse, but otherwise we're just duplicating code
# HelpfulSubparserAction gives more information about the possible choices from a subparsed choice
# HelpfulArgParser gives the list of choices when no arguments are provided to a choice option whilst still using a METAVAR
class HelpfulSubparserAction(argparse._SubParsersAction):
"""Class to either select a unique plugin based on a substring, or identify
the alternatives."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# We don't want the action self-check to kick in, so we remove the choices list, the check happens in __call__
self.choices = None
def __call__(self, parser, namespace, values, option_string = None):
parser_name = values[0]
arg_strings = values[1:]
# set the parser name if requested
if self.dest is not argparse.SUPPRESS:
setattr(namespace, self.dest, parser_name)
matched_parsers = [name for name in self._name_parser_map if parser_name in name]
if len(matched_parsers) < 1:
msg = 'invalid choice {} (choose from {})'.format(parser_name, ', '.join(self._name_parser_map))
raise argparse.ArgumentError(self, msg)
if len(matched_parsers) > 1:
msg = 'plugin {} matches multiple plugins ({})'.format(parser_name, ', '.join(matched_parsers))
raise argparse.ArgumentError(self, msg)
parser = self._name_parser_map[matched_parsers[0]]
setattr(namespace, 'plugin', matched_parsers[0])
# parse all the remaining options into the namespace
# store any unrecognized options on the object, so that the top
# level parser can decide what to do with them
# In case this subparser defines new defaults, we parse them
# in a new namespace object and then update the original
# namespace for the relevant parts.
subnamespace, arg_strings = parser.parse_known_args(arg_strings, None)
for key, value in vars(subnamespace).items():
setattr(namespace, key, value)
if arg_strings:
vars(namespace).setdefault(argparse._UNRECOGNIZED_ARGS_ATTR, [])
getattr(namespace, argparse._UNRECOGNIZED_ARGS_ATTR).extend(arg_strings)
class HelpfulArgParser(argparse.ArgumentParser):
def _match_argument(self, action, arg_strings_pattern) -> int:
# match the pattern for this action to the arg strings
nargs_pattern = self._get_nargs_pattern(action)
match = re.match(nargs_pattern, arg_strings_pattern)
# raise an exception if we weren't able to find a match
if match is None:
nargs_errors = {
None: gettext.gettext('expected one argument'),
argparse.OPTIONAL: gettext.gettext('expected at most one argument'),
argparse.ONE_OR_MORE: gettext.gettext('expected at least one argument'),
}
msg = nargs_errors.get(action.nargs)
if msg is None:
msg = gettext.ngettext('expected %s argument', 'expected %s arguments', action.nargs) % action.nargs
if action.choices:
msg = "{} (from: {})".format(msg, ", ".join(action.choices))
raise argparse.ArgumentError(action, msg)
# return the number of arguments matched
return len(match.group(1))