From b647c3a089be567090f2050cf3c9ffe0f5c539b0 Mon Sep 17 00:00:00 2001
From: Mike Auty <mike.auty@gmail.com>
Date: Wed, 30 May 2018 22:36:42 +0100
Subject: [PATCH] Refactor volshell with inheritance and add dt.

It makes enourmous sense to inherit the windows volshell plugin from the
volshell one.  I also added the generic dt function.
---
 volatility/cli/volshell/shellplugin.py | 32 +++++++++++++++++++++++---
 volatility/cli/volshell/windows.py     | 18 ++++++++++-----
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/volatility/cli/volshell/shellplugin.py b/volatility/cli/volshell/shellplugin.py
index cea6f631f..829b729a5 100644
--- a/volatility/cli/volshell/shellplugin.py
+++ b/volatility/cli/volshell/shellplugin.py
@@ -1,12 +1,12 @@
 import code
 import inspect
+import typing
 
-from volatility.framework import renderers
+from volatility.framework import renderers, interfaces
 from volatility.framework.configuration import requirements
-from volatility.framework.interfaces import plugins
 
 
-class Volshell(plugins.PluginInterface):
+class Volshell(interfaces.plugins.PluginInterface):
     """Shell environment to directly interact with a memory image"""
 
     @classmethod
@@ -31,6 +31,8 @@ class Volshell(plugins.PluginInterface):
         if additional_locals is not None:
             vars.update(additional_locals)
 
+        vars.update(self.load_functions())
+
         # Try to enable tab completion
         try:
             import readline
@@ -48,3 +50,27 @@ class Volshell(plugins.PluginInterface):
         code.interact(local = vars)
 
         return renderers.TreeGrid([], lambda: [])
+
+    def load_functions(self) -> typing.Dict[str, typing.Callable]:
+        """Returns a dictionary listing the functions to be added to the environment"""
+        return {"dt": self.display_type}
+
+    def display_type(self, object: interfaces.objects.ObjectInterface):
+        """Display Type"""
+        longest_member = longest_offset = 0
+        for member in object.vol.members:
+            relative_offset, member_type = object.vol.members[member]
+            longest_member = max(len(member), longest_member)
+            longest_offset = max(len(hex(relative_offset)), longest_offset)
+
+        for member in object.vol.members:
+            relative_offset, member_type = object.vol.members[member]
+            len_offset = len(hex(relative_offset))
+            len_member = len(member)
+            print(" " * (longest_offset - len_offset),
+                  hex(relative_offset),
+                  "\t\t",
+                  member,
+                  " " * (longest_member - len_member),
+                  "\t\t",
+                  member_type.vol.type_name)
diff --git a/volatility/cli/volshell/windows.py b/volatility/cli/volshell/windows.py
index 099568496..b6afd782b 100644
--- a/volatility/cli/volshell/windows.py
+++ b/volatility/cli/volshell/windows.py
@@ -1,16 +1,16 @@
 import inspect
+import typing
 
 from volatility.cli.volshell import shellplugin
 from volatility.framework.configuration import requirements
-from volatility.framework.interfaces import plugins
 
 
-class Volshell(plugins.PluginInterface):
+class Volshell(shellplugin.Volshell):
     """Shell environment to directly interact with a windows memory image"""
 
     @classmethod
     def get_requirements(cls):
-        return (shellplugin.Volshell.get_requirements() +
+        return (super().get_requirements() +
                 [requirements.SymbolRequirement(name = "nt_symbols", description = "Windows OS"),
                  requirements.IntRequirement(name = 'pid',
                                              description = "Process ID",
@@ -44,7 +44,14 @@ class Volshell(plugins.PluginInterface):
         for proc in eproc.ActiveProcessLinks:
             yield proc
 
-    def run(self):
+    def load_functions(self) -> typing.Dict[str, typing.Callable]:
+        result = super().load_functions()
+        result.update({
+            'ps': lambda: list(self.list_processes())
+        })
+        return result
+
+    def run(self, additional_locals = None):
         # Determine locals
         curframe = inspect.currentframe()
 
@@ -52,7 +59,6 @@ class Volshell(plugins.PluginInterface):
         layer_name = self.config['primary']
         kvo = self.context.memory[layer_name].config['kernel_virtual_offset']
         nt = self.context.module(self.config['nt_symbols'], layer_name = layer_name, offset = kvo)
-
         ps = lambda: list(self.list_processes())
 
         pid = self.config.get('pid', None)
@@ -63,4 +69,4 @@ class Volshell(plugins.PluginInterface):
                     eproc = _x
                     break
 
-        return shellplugin.Volshell(self.context, self.config_path).run(curframe.f_locals)
+        return super().run(curframe.f_locals)