diff --git a/volatility3/cli/volshell/generic.py b/volatility3/cli/volshell/generic.py index 2321408fe..143e26500 100644 --- a/volatility3/cli/volshell/generic.py +++ b/volatility3/cli/volshell/generic.py @@ -8,6 +8,7 @@ import random import string import struct import sys +import textwrap from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union from urllib import parse, request @@ -23,6 +24,14 @@ try: except ImportError: has_capstone = False +try: + from IPython import terminal + from traitlets import config as traitlets_config + + has_ipython = True +except ImportError: + has_ipython = False + class Volshell(interfaces.plugins.PluginInterface): """Shell environment to directly interact with a memory image.""" @@ -69,43 +78,67 @@ class Volshell(interfaces.plugins.PluginInterface): """ # Try to enable tab completion - try: - import readline - except ImportError: - pass - else: - import rlcompleter + if not has_ipython: + try: + import readline + import rlcompleter - completer = rlcompleter.Completer(namespace=self._construct_locals_dict()) - readline.set_completer(completer.complete) - readline.parse_and_bind("tab: complete") - print("Readline imported successfully") + completer = rlcompleter.Completer( + namespace=self._construct_locals_dict() + ) + readline.set_completer(completer.complete) + readline.parse_and_bind("tab: complete") + print("Readline imported successfully") + except ImportError: + print( + "Readline or rlcompleter module could not be imported. Tab completion will not be available." + ) # TODO: provide help, consider generic functions (pslist?) and/or providing windows/linux functions mode = self.__module__.split(".")[-1] mode = mode[0].upper() + mode[1:] - banner = f""" - Call help() to see available functions + banner = textwrap.dedent( + f""" + Call help() to see available functions - Volshell mode : {mode} - Current Layer : {self.current_layer} - Current Symbol Table : {self.current_symbol_table} - Current Kernel Name : {self.current_kernel_name} -""" + Volshell mode : {mode} + Current Layer : {self.current_layer} + Current Symbol Table : {self.current_symbol_table} + Current Kernel Name : {self.current_kernel_name} + """ + ) sys.ps1 = f"({self.current_layer}) >>> " # Dict self._construct_locals_dict() will have priority on keys combined_locals = additional_locals.copy() combined_locals.update(self._construct_locals_dict()) - self.__console = code.InteractiveConsole(locals=combined_locals) + if has_ipython: + + class LayerNamePrompt(terminal.prompts.Prompts): + def in_prompt_tokens(self, cli=None): + slf = self.shell.user_ns.get("self") + layer_name = slf.current_layer if slf else "no_layer" + return [(terminal.prompts.Token.Prompt, f"[{layer_name}]> ")] + + c = traitlets_config.Config() + c.TerminalInteractiveShell.prompts_class = LayerNamePrompt + c.InteractiveShellEmbed.banner2 = banner + self.__console = terminal.embed.InteractiveShellEmbed( + config=c, user_ns=combined_locals + ) + else: + self.__console = code.InteractiveConsole(locals=combined_locals) # Since we have to do work to add the option only once for all different modes of volshell, we can't # rely on the default having been set if self.config.get("script", None) is not None: self.run_script(location=self.config["script"]) - self.__console.interact(banner=banner) + if has_ipython: + self.__console() + else: + self.__console.interact(banner=banner) return renderers.TreeGrid([("Terminating", str)], None) @@ -508,10 +541,13 @@ class Volshell(interfaces.plugins.PluginInterface): location = "file:" + request.pathname2url(location) print(f"Running code from {location}\n") accessor = resources.ResourceAccessor() - with accessor.open(url=location) as fp: - self.__console.runsource( - io.TextIOWrapper(fp, encoding="utf-8").read(), symbol="exec" - ) + with accessor.open(url=location) as handle, io.TextIOWrapper( + handle, encoding="utf-8" + ) as fp: + if has_ipython: + self.__console.ex(fp.read()) + else: + self.__console.runsource(fp.read(), symbol="exec") print("\nCode complete") def load_file(self, location: str):