From 44f26c928eddaaf6743ac07b22331ad082e86bc1 Mon Sep 17 00:00:00 2001
From: Mike Auty <mike.auty@gmail.com>
Date: Sun, 28 Jul 2024 12:35:02 +0100
Subject: [PATCH] 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)