From 44f26c928eddaaf6743ac07b22331ad082e86bc1 Mon Sep 17 00:00:00 2001 From: Mike Auty Date: Sun, 28 Jul 2024 12:35:02 +0100 Subject: [PATCH 1/6] Add in shtab autocompletion --- volatility3/cli/__init__.py | 22 +++++++++++++++++++--- volatility3/cli/volargparse.py | 6 ++++-- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/volatility3/cli/__init__.py b/volatility3/cli/__init__.py index 6b17edac0..883b54ac4 100644 --- a/volatility3/cli/__init__.py +++ b/volatility3/cli/__init__.py @@ -22,6 +22,13 @@ import traceback from typing import Any, Dict, List, Tuple, Type, Union from urllib import parse, request +try: + import shtab + + HAS_SHTAB = True +except ImportError: + HAS_SHTAB = False + from volatility3.cli import text_filter import volatility3.plugins import volatility3.symbols @@ -106,6 +113,9 @@ class CommandLine: ] ) + # Argument for doing autocompletion + print_completion_arg = "--print-completion" + # Load up system defaults delayed_logs, default_config = self.load_system_defaults("vol.json") @@ -246,10 +256,12 @@ class CommandLine: # We have to filter out help, otherwise parse_known_args will trigger the help message before having # processed the plugin choice or had the plugin subparser added. known_args = [arg for arg in sys.argv if arg != "--help" and arg != "-h"] - partial_args, _ = parser.parse_known_args(known_args) - + partial_args, unknown_args = parser.parse_known_args(known_args) banner_output = sys.stdout - if renderers[partial_args.renderer].structured_output: + if ( + renderers[partial_args.renderer].structured_output + or print_completion_arg in unknown_args + ): banner_output = sys.stderr banner_output.write(f"Volatility 3 Framework {constants.PACKAGE_VERSION}\n") @@ -351,6 +363,10 @@ class CommandLine: # Hand the plugin requirements over to the CLI (us) and let it construct the config tree # Run the argparser + if HAS_SHTAB: + # The autocompletion line must be after the partial_arg handling, so that it doesn't trip it + # before all the plugins have been added + shtab.add_argument_to(parser, [print_completion_arg]) args = parser.parse_args() if args.plugin is None: parser.error("Please select a plugin to run") diff --git a/volatility3/cli/volargparse.py b/volatility3/cli/volargparse.py index 3048a0885..3e7eb6751 100644 --- a/volatility3/cli/volargparse.py +++ b/volatility3/cli/volargparse.py @@ -21,8 +21,6 @@ class HelpfulSubparserAction(argparse._SubParsersAction): def __init__(self, *args, **kwargs) -> None: 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, @@ -100,3 +98,7 @@ class HelpfulArgParser(argparse.ArgumentParser): # return the number of arguments matched return len(match.group(1)) + + def _check_value(self, action, value): + if not isinstance(action, HelpfulSubparserAction): + return super()._check_value(action, value) From e89e77637776b14c61516d2c47a7148a2f13860a Mon Sep 17 00:00:00 2001 From: Mike Auty Date: Sun, 28 Jul 2024 12:41:31 +0100 Subject: [PATCH 2/6] Try out argcomplete as well --- vol.py | 1 + volatility3/cli/__init__.py | 21 ++++++++------------- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/vol.py b/vol.py index ff420cad5..c49d5985d 100755 --- a/vol.py +++ b/vol.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +# PYTHON_ARGCOMPLETE_OK # This file is Copyright 2019 Volatility Foundation and licensed under the Volatility Software License 1.0 # which is available at https://www.volatilityfoundation.org/license/vsl-v1.0 diff --git a/volatility3/cli/__init__.py b/volatility3/cli/__init__.py index 883b54ac4..1209d7cdc 100644 --- a/volatility3/cli/__init__.py +++ b/volatility3/cli/__init__.py @@ -23,11 +23,11 @@ from typing import Any, Dict, List, Tuple, Type, Union from urllib import parse, request try: - import shtab + import argcomplete - HAS_SHTAB = True + HAS_ARGCOMPLETE = True except ImportError: - HAS_SHTAB = False + HAS_ARGCOMPLETE = False from volatility3.cli import text_filter import volatility3.plugins @@ -113,9 +113,6 @@ class CommandLine: ] ) - # Argument for doing autocompletion - print_completion_arg = "--print-completion" - # Load up system defaults delayed_logs, default_config = self.load_system_defaults("vol.json") @@ -256,12 +253,10 @@ class CommandLine: # We have to filter out help, otherwise parse_known_args will trigger the help message before having # processed the plugin choice or had the plugin subparser added. known_args = [arg for arg in sys.argv if arg != "--help" and arg != "-h"] - partial_args, unknown_args = parser.parse_known_args(known_args) + partial_args, _ = parser.parse_known_args(known_args) + banner_output = sys.stdout - if ( - renderers[partial_args.renderer].structured_output - or print_completion_arg in unknown_args - ): + if renderers[partial_args.renderer].structured_output: banner_output = sys.stderr banner_output.write(f"Volatility 3 Framework {constants.PACKAGE_VERSION}\n") @@ -363,10 +358,10 @@ class CommandLine: # Hand the plugin requirements over to the CLI (us) and let it construct the config tree # Run the argparser - if HAS_SHTAB: + if HAS_ARGCOMPLETE: # The autocompletion line must be after the partial_arg handling, so that it doesn't trip it # before all the plugins have been added - shtab.add_argument_to(parser, [print_completion_arg]) + argcomplete.autocomplete(parser) args = parser.parse_args() if args.plugin is None: parser.error("Please select a plugin to run") From 73bc10c2834d9a8fa2ab04c098a4735542226170 Mon Sep 17 00:00:00 2001 From: Mike Auty Date: Wed, 31 Jul 2024 21:46:57 +0100 Subject: [PATCH 3/6] Wire the argcomplete into volshell too --- volatility3/cli/volshell/__init__.py | 12 ++++++++++++ volshell.py | 1 + 2 files changed, 13 insertions(+) diff --git a/volatility3/cli/volshell/__init__.py b/volatility3/cli/volshell/__init__.py index 035ed9b2e..2bf1958e2 100644 --- a/volatility3/cli/volshell/__init__.py +++ b/volatility3/cli/volshell/__init__.py @@ -21,6 +21,14 @@ from volatility3.framework import ( plugins, ) +try: + import argcomplete + + HAS_ARGCOMPLETE = True +except ImportError: + HAS_ARGCOMPLETE = False + + # Make sure we log everything rootlog = logging.getLogger() @@ -276,6 +284,10 @@ class VolShell(cli.CommandLine): # Hand the plugin requirements over to the CLI (us) and let it construct the config tree # Run the argparser + if HAS_ARGCOMPLETE: + # The autocompletion line must be after the partial_arg handling, so that it doesn't trip it + # before all the plugins have been added + argcomplete.autocomplete(parser) args = parser.parse_args() vollog.log( diff --git a/volshell.py b/volshell.py index 71d35a47c..65b11885e 100755 --- a/volshell.py +++ b/volshell.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +# PYTHON_ARGCOMPLETE_OK # This file is Copyright 2019 Volatility Foundation and licensed under the Volatility Software License 1.0 # which is available at https://www.volatilityfoundation.org/license/vsl-v1.0 From d576f8cb48d064ce3dc87936df642481aa1f3186 Mon Sep 17 00:00:00 2001 From: Mike Auty Date: Wed, 31 Jul 2024 21:49:07 +0100 Subject: [PATCH 4/6] Fix CodeQL error --- volatility3/cli/volargparse.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/volatility3/cli/volargparse.py b/volatility3/cli/volargparse.py index 3e7eb6751..2bd53077b 100644 --- a/volatility3/cli/volargparse.py +++ b/volatility3/cli/volargparse.py @@ -99,6 +99,7 @@ class HelpfulArgParser(argparse.ArgumentParser): # return the number of arguments matched return len(match.group(1)) - def _check_value(self, action, value): + def _check_value(self, action: argparse.Action, value: Any) -> None: if not isinstance(action, HelpfulSubparserAction): return super()._check_value(action, value) + return None From e6308a6035156cab5abb6f0cdf537fb75e5881e8 Mon Sep 17 00:00:00 2001 From: Mike Auty Date: Thu, 1 Aug 2024 21:04:24 +0100 Subject: [PATCH 5/6] Make suggested changes by gcmoreira --- volatility3/cli/volargparse.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/volatility3/cli/volargparse.py b/volatility3/cli/volargparse.py index 2bd53077b..5ce2646ed 100644 --- a/volatility3/cli/volargparse.py +++ b/volatility3/cli/volargparse.py @@ -100,6 +100,11 @@ class HelpfulArgParser(argparse.ArgumentParser): return len(match.group(1)) def _check_value(self, action: argparse.Action, value: Any) -> None: + """This is called to ensure a value is correct/valid + This fails when we want to accept partial values for the plugin name, + so we disable the check (which will throw ArgumentErrors for failed checks) + but only for our plugin subparser, so all other arguments are checked correctly + """ if not isinstance(action, HelpfulSubparserAction): - return super()._check_value(action, value) + super()._check_value(action, value) return None From 0bda1543854d8f92e7fb2bf4c5eff3afdf117819 Mon Sep 17 00:00:00 2001 From: Mike Auty Date: Thu, 1 Aug 2024 21:08:18 +0100 Subject: [PATCH 6/6] Clarify the documentation a little --- volatility3/cli/volargparse.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/volatility3/cli/volargparse.py b/volatility3/cli/volargparse.py index 5ce2646ed..fd61ddce0 100644 --- a/volatility3/cli/volargparse.py +++ b/volatility3/cli/volargparse.py @@ -101,9 +101,16 @@ class HelpfulArgParser(argparse.ArgumentParser): def _check_value(self, action: argparse.Action, value: Any) -> None: """This is called to ensure a value is correct/valid - This fails when we want to accept partial values for the plugin name, - so we disable the check (which will throw ArgumentErrors for failed checks) - but only for our plugin subparser, so all other arguments are checked correctly + + In normal operation, it would check that a value provided is valid and return None + If it was not valid, it would throw an ArgumentError + + When people provide a partial plugin name, we want to look for a matching plugin name + which happens in the HelpfulSubparserAction's __call_method + + To get there without tripping the check_value failure, we have to prevent the exception + being thrown when the value is a HelpfulSubparserAction. This therefore affects no other + checks for normal parameters. """ if not isinstance(action, HelpfulSubparserAction): super()._check_value(action, value)