Add in file producer/consumer API.

volatility/cli/__init__.py

@@ -12,6 +12,7 @@ import argparse
 import inspect
 import json
 import logging
+import os
 import sys
 import typing
 
@@ -53,11 +54,11 @@ class PrintedProgress(object):
         print(message, end = ' ' * (self._max_message_len - message_len))
 
 
-class CommandLine(object):
+class CommandLine(interfaces.plugins.FileConsumerInterface):
     """Constructs a command-line interface object for users to run plugins"""
 
     def __init__(self):
-        pass
+        self.output_dir = None
 
     def run(self):
         """Executes the command line module, taking the system arguments, determining the plugin to run and then running it"""
@@ -75,6 +76,8 @@ class CommandLine(object):
         parser.add_argument("-p", "--plugins", help = "Semi-colon separated list of paths to find plugins",
                             default = "", type = str)
         parser.add_argument("-v", "--verbosity", help = "Increase output verbosity", default = 0, action = "count")
+        parser.add_argument("-o", "--output-dir", help = "Directory in which to output any generated files",
+                            default = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')), type = str)
         parser.add_argument("-q", "--quiet", help = "Remove progress feedback", default = False, action = 'store_true')
         parser.add_argument("-l", "--log", help = "Log output to a file as well as the console", default = None,
                             type = str)
@@ -204,9 +207,29 @@ class CommandLine(object):
             with open("config.json", "w") as f:
                 json.dump(dict(constructed.build_configuration()), f, sort_keys = True, indent = 2)
 
+        self.output_dir = args.output_dir
+        constructed.set_file_consumer(self)
+
         # Construct and run the plugin
         text.QuickTextRenderer().render(constructed.run())
 
+    def consume_file(self, filedata: interfaces.plugins.FileInterface):
+        """Consumes a file as produced by a plugin"""
+        if self.output_dir is None:
+            raise ValueError("Output directory has not been correctly specified")
+        os.makedirs(self.output_dir, exist_ok = True)
+
+        pref_name_array = filedata.preferred_filename.split('.')
+        filename, extension = os.path.join(self.output_dir, '.'.join(pref_name_array[:-1])), pref_name_array[-1]
+        output_filename = "{}.{}".format(filename, extension)
+
+        if not os.path.exists(output_filename):
+            with open(output_filename, "wb") as current_file:
+                current_file.write(filedata.data.getbuffer())
+                vollog.log(logging.INFO, "Saved stored plugin file: {}".format(output_filename))
+        else:
+            vollog.warning("Refusing to overwrite an existing file: {}".format(output_filename))
+
     def populate_requirements_argparse(self,
                                        parser: argparse.ArgumentParser,
                                        configurable: typing.Type[interfaces.configuration.ConfigurableInterface]):

volatility/framework/interfaces/plugins.py

@@ -4,6 +4,7 @@ They are called and carry out some algorithms on data stored in layers using obj
 """
 
 # Configuration interfaces must be imported separately, since we're part of interfaces and can't import ourselves
+import io
 import logging
 import typing
 from abc import ABCMeta, abstractmethod
@@ -18,6 +19,25 @@ if typing.TYPE_CHECKING:
     from volatility.framework import interfaces, renderers
 
 
+class FileInterface(validity.ValidityRoutines, metaclass = ABCMeta):
+    """Class for storing Files in the plugin as a means to output a file or files when necessary"""
+
+    def __init__(self, filename: str, data: bytes = None):
+        self.preferred_filename = filename
+        self.data = io.BytesIO(data)
+
+
+class FileConsumerInterface(object):
+    """Class for consuming files potentially produced by plugins
+
+    We use the producer/consumer model to ensure we can avoid running out of memory by storing every file produced.
+    The downside is, we can't provide much feedback to the producer about what happened to their file (other than exceptions).
+    """
+
+    def consume_file(self, file: FileInterface) -> None:
+        """Consumes a file as passed back to a UI by a plugin"""
+
+
 #
 # Plugins
 # - Take in relevant number of TranslationLayers (of specified type)
@@ -49,6 +69,17 @@ class PluginInterface(interfaces_configuration.ConfigurableInterface, validity.V
         if self.unsatisfied(context, config_path):
             vollog.warning("Plugin failed validation")
             raise exceptions.PluginRequirementException("The plugin configuration failed to validate")
+        self._file_consumer = None  # type: FileConsumerInterface
+
+    def set_file_consumer(self, consumer: FileConsumerInterface) -> None:
+        self._file_consumer = self._check_type(consumer, FileConsumerInterface)
+
+    def produce_file(self, filedata: FileInterface) -> None:
+        """Adds a file to the plugin's file store and returns the chosen filename for the file"""
+        if self._file_consumer:
+            self._file_consumer.consume_file(filedata)
+        else:
+            vollog.debug("No file consumer specified to consume: {}".format(filedata.preferred_filename))
 
     @classmethod
     def get_requirements(cls) -> typing.List['interfaces.configuration.RequirementInterface']:

volatility/plugins/windows/vaddump.py

@@ -1,24 +1,21 @@
-import os
-import volatility.framework.interfaces.plugins as interfaces_plugins
-import volatility.plugins.windows.vadinfo as vadinfo
-import volatility.plugins.windows.pslist as pslist
-from volatility.framework import renderers
-from volatility.framework.objects import utility
-from volatility.framework.configuration import requirements
 import logging
 
+import volatility.framework.interfaces.plugins as interfaces_plugins
+import volatility.plugins.windows.pslist as pslist
+import volatility.plugins.windows.vadinfo as vadinfo
+from volatility.framework import renderers
+from volatility.framework.objects import utility
+
 vollog = logging.getLogger()
 
+
 class VadDump(interfaces_plugins.PluginInterface):
     """Dumps process memory ranges"""
 
     @classmethod
     def get_requirements(cls):
         # Since we're calling the plugin, make sure we have the plugin's requirements
-        return vadinfo.VadInfo.get_requirements() + [requirements.StringRequirement(name = "outdir",
-                                               description = "Output directory",
-                                               default = None,
-                                               optional = False)]
+        return vadinfo.VadInfo.get_requirements()
 
     def _generator(self, procs):
 
@@ -34,26 +31,23 @@ class VadDump(interfaces_plugins.PluginInterface):
 
             for vad in plugin.list_vads(proc):
                 try:
-                    file_name = os.path.join(self.config["outdir"],
-                                             "pid.{0}.vad.{1:#x}-{2:#x}.dmp".format(proc.UniqueProcessId, vad.get_start(), vad.get_end()))
+                    filedata = interfaces_plugins.FileInterface(
+                        "pid.{0}.vad.{1:#x}-{2:#x}.dmp".format(proc.UniqueProcessId,
+                                                               vad.get_start(),
+                                                               vad.get_end()))
 
-                    if os.path.exists(file_name):
-                        raise FileExistsError
+                    offset = vad.get_start()
+                    out_of_range = vad.get_start() + vad.get_end()
+                    while offset < out_of_range:
+                        to_read = min(chunk_size, out_of_range - offset)
+                        data = proc_layer.read(offset, to_read, pad = True)
+                        if not data:
+                            break
+                        filedata.data.write(data)
+                        offset += to_read
 
-                    with open(file_name, "wb") as handle:
-                        offset = vad.get_start()
-                        out_of_range = vad.get_start() + vad.get_end()
-                        while offset < out_of_range:
-                            to_read = min(chunk_size, out_of_range - offset)
-                            data = proc_layer.read(offset, to_read, pad = True)
-                            if not data:
-                                break
-                            handle.write(data)
-                            offset += to_read
-
-                    result_text = "Saved to {}".format(os.path.basename(file_name))
-                except FileExistsError:
-                    result_text = "Refusing to overwrite the existing {}".format(file_name)
+                    self.produce_file(filedata)
+                    result_text = "Stored {}".format(filedata.preferred_filename)
                 except Exception:
                     result_text = "Unable to dump {0:#x} - {1:#x}".format(vad.get_start(), vad.get_end())
 
@@ -62,16 +56,6 @@ class VadDump(interfaces_plugins.PluginInterface):
                            result_text))
 
     def run(self):
-
-        try:
-            # the optional=False requirement should make sure this always exists
-            os.makedirs(self.config["outdir"])
-        except FileExistsError:
-            pass
-        except OSError:
-            # is this what we want to raise here?
-            raise OSError("Cannot create the desired output directory!")
-
         plugin = pslist.PsList(self.context, "plugins.VadDump")
 
         return renderers.TreeGrid([("PID", int),