Merge pull request #1229 from volatilityfoundation/feature/argcomplete

Feature/argcomplete

Author: ikelos

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

volatility3/cli/__init__.py

@@ -22,6 +22,13 @@
 from typing import Any, Dict, List, Tuple, Type, Union
 from urllib import parse, request
 
+try:
+    import argcomplete
+
+    HAS_ARGCOMPLETE = True
+except ImportError:
+    HAS_ARGCOMPLETE = False
+
 from volatility3.cli import text_filter
 import volatility3.plugins
 import volatility3.symbols
@@ -351,6 +358,10 @@ def run(self):
         # 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()
         if args.plugin is None:
             parser.error("Please select a plugin to run")

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,20 @@ def _match_argument(self, action, arg_strings_pattern) -> int:
 
         # return the number of arguments matched
         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
+
+        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)
+        return None

volatility3/cli/volshell/__init__.py

@@ -21,6 +21,14 @@
     plugins,
 )
 
+try:
+    import argcomplete
+
+    HAS_ARGCOMPLETE = True
+except ImportError:
+    HAS_ARGCOMPLETE = False
+
+
 # Make sure we log everything
 
 rootlog = logging.getLogger()
@@ -276,6 +284,10 @@ def run(self):
         # 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(

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