volatilityfoundation-volatility/volatility/cache.py

# This file is part of Volatility.
#
# Volatility is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# Volatility is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Volatility.  If not, see <http://www.gnu.org/licenses/>.

""" This module implements the volatility caching subsystem.

The volatility caching subsystem has the following design goals:

 1) Ability to cache arbitrary objects - The allows complex objects to
    be cached for later retrieval. For example, objects may be as
    simple as constants for KPCR addresses, to entire x86 page
    translation tables, or even hibernation decompression
    datastructures. To achieve this we use the standard python pickle
    system. In many use cases, the cache needs to facilitate
    persistant memoising of functions and generators (more on that
    below).

 2) Cached objects are stored by a hierarchical key namespace. Keys
    are specified in a URL notation. By default, relative URLs are
    interpreted relative to the memory image location (the value of
    the --location option). This scheme allows us to specify both
    global (per installation) and per image keys. For example given an
    image located in /tmp/foobar.img:

    - file:///tmp/foobar.img/kernel/debugging/KPCR refers to this
      image's KPCR location.

    - file:///tmp/foobar.img/address_spaces/memory_translation/pdpte
      refers to the cached page tables.

    - http://www.volatility.org/schema#configuration/renderer specifies
      the currently configured renderer (i.e. its a global setting).

 3) Storage of the cache is abstracted and selectable via the
 --cache_engine configuration variable. This allows the separation
 from the concerete storage of the cache and the abstraction of the
 cache in a running process.

Abstraction of Cache
--------------------

Within the running volatiltiy framework the cache appears as an
abstract tree with nodes inherited from the CacheNode class:

class CacheNode(object):
    def __init__(self, name, parent, payload = None):
        ''' Creates a new Cache node under the parent. The new node
        will carry the specified payload
        '''

    def __str__(self):
        ''' Produce a human readable version of the payload '''

    def set_payload(self, payload):
        ''' Update the current payload with the new specified payload '''

    def dump(self):
        ''' Dump the node to disk for later retrieval. This is
        normally called when the process has exited. '''

    def get_payload(self):
       ''' retrieve this node's payload '''

In order to check the cache, plugins issue the Cache.Check() function:

def Check(path, callback = None, cls = CacheNode):
    ''' Traverse the cache tree and retrieve the stored CacheNode.

    If there is no such stored CacheNode and callback is specified,
    attempt to create it using the cache_node_class with the payload
    returned from the callback. If callback is not specified we just
    return None.

Decorators
----------
You can also use the cache decorator to cache the results of any
function - this is probably the easiest way to apply caching to
existing code. For example, suppose we want to cache the results of
the psscan plugin:

class PSScan(commands.Command):
....
   @cache("/scanners/psscan")
   def calculate(self):
       .....

This will automatically create the CacheNode at the specified tree
location (note that since the URL is given as a relative URL it is
based at the current value of the --location - that means it applies
to the current memory image only).

Note that since calculate() returns a generator, the decorator will
also return a generator - It will not iterate over the calculate
method unnecessarily, but will yield results immediately. This does
not compromise performance in the case of a cache miss. Unfortunately
this also means that if the generator is stopped prematurely, we are
unable to cache the result set in the general case. This is the only
caveat on caching generators.

Storage classes
---------------
The cache system discussed above can be thought of as an abstract
construct in the process memory. To make it persistant on disk we have
the storage class (which can be selected using the --cache_engine
directive). The following cache engines are implemented:

File Storage
============
This is the default cache engine. We simply maintain a directory
structure which corresponds to the URL of the key after applying the
appropriate filesystem safe escaping operation. Objects are stored in
stand alone files using the pickle module.

Zip Storage
===========
This storage is essentially the same as the File storage above, except
that the cache directory for each image file is maintained in a Zip
file stored at the --cache_direcory directive with the same filename
as the image and a .zip extension.


Use cases
---------
The following common use cases are discussed:

1) Dynamic address spaces. In some address spaces memory address
   mappings can not be cached since they change all the time. For
   example in the firewire address space, it is incorrect to cache any
   page translations or scanning results etc. This is easily achieved
   by having the firewire address space store a BlockingCacheNode()
   instance at critical tree nodes. These prevent new nodes from being
   inserted into the tree and force a cache miss whenever any keys are
   searched under these nodes. Note that this still allows the cache
   to store the locations of things which might not change, even for
   live memory analysis, such as KPCR locations.

2) History logging and audit logs. Currently volatility works by
   running the framework multiple times on the same plugin with
   different command line options. This can be audited using the
   caching system by storing the current command line in a specific
   location using a specific CacheNode. This implementation can be
   used to append new commandlines to the same key. Configuration
   options can also become sticky in this way and remember the same
   values they had previously. This avoid users having to append many
   command line arguements (i.e. having to specify --profile, --kpcr,
   --dtb on every command line).

3) Unit tests.  Unit tests can be easily implemented using the caching
   subsystem as follows:

   - A test() method is added to each plugin. Usually this is actually
     the same as calculate().

   - This method is decorated to be cached under the
     "/tests/pluginname" key (i.e. relative to the current image). The
     CacheNode implementation is TestCacheNode which implements a
     special update_payload() method. The TestCacheNode also ensures
     that cache miss always occurs (by implementing a get_payload()
     method which returns None).

   - The update_payload() method ensures that the old payload and the
     new payloads are the same (if they are generators we ensure each
     member is the same as well - using the __eq__ method).

   The overall result is that unit tests can be run on any image as
   normal. If the particular test was never run on the image, we just
   cache the result of the plugin. If on the other hand, the result
   was already run on this image, the old result is compared to the
   new result and if a discrepancy is detected, an exception is
   raised.

   This testing framework is easy to implement and automatically
   guards against regression bugs. Since we use the __eq__ method of
   arbitrary objects, its also not limited to testing text string
   matches. For example, the object framework defines two objects are
   being equal if they are of the same type and they point at the same
   address. Even if the textual representation of the object's
   printouts has changed between versions, as long as the same objects
   are found in both cases no regressions will be reported.

4) Reporting framework. By having a persistant caching framework we
   now have the concept of a volatility analysis session. In other
   words, each new execution of volatility adds new information to
   what we know about the image. This new information is stored in the
   cache tree. We can actually produce a full report from the cache
   tree by traversing all the CacheNodes and calling their __str__()
   methods.

   If caching is introduced via decorators, the CacheNode already
   knows about the render() method of the plugin and can automatically
   generate the output from the plugin (this is very fast as the
   calculate is received from the cache). We therefore can generate a
   full report of all the plugins very quickly automatically.

   By default CacheNodes have an empty __str__() methods, so things
   like pas2kas lookup tables are not reported. Specialised reporting
   functions can be made if needed by implementing __str__() functions
   as needed.

"""
import types
import os
import urlparse
import volatility.conf as conf
import volatility.obj as obj
import volatility.debug as debug
import volatility.exceptions as exceptions
import cPickle as pickle
config = conf.ConfObject()

## Where to stick the cache
default_cache_location = os.path.join((os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")), "volatility")

config.add_option("CACHE-DIRECTORY", default = default_cache_location,
                  cache_invalidator = False,
                  help = "Directory where cache files are stored")

class CacheContainsGenerator(exceptions.VolatilityException):
    """Exception raised when the cache contains a generator"""
    pass

class InvalidCache(Exception):
    """Exception raised when the cache item is determined to be invalid."""
    pass

class CacheNode(object):
    """ Base class for Cache nodes """
    def __init__(self, name, stem, storage = None, payload = None, invalidator = None):
        ''' Creates a new Cache node under the parent. The new node
        will carry the specified payload
        '''
        self.name = name
        self.payload = payload
        self.storage = storage
        self.stem = stem

        # This object encapsulate the running environment. If the
        # environment during the time of unpickling differs from the
        # environment during the time of pickling we refuse to
        # unpickle this object, and the cache misses. We dont really
        # do anything with it, just have it serialised as well.
        self.invalidator = invalidator

    def __getitem__(self, item = ''):
        item_url = "{0}/{1}".format(self.stem, item)

        ## Try to load it from the storage manager
        try:
            result = self.storage.load(item_url)
            if result:
                return result
        except Exception, e:
            raise KeyError(e)

        ## Make a new empty Node instead on demand
        raise KeyError("item not found")

    def __str__(self):
        ''' Produce a human readable version of the payload. '''
        return ''

    def _find_generators(self, item):
        """ A recursive function to flatten generators into lists """
        try:
            result = []
            # Make sure dicts aren't flattened to lists
            if isinstance(item, dict):
                result = {}
                for i in item:
                    result[self._find_generators(i)] = self._find_generators(item[i])
                return result

            # Since NoneObjects and strings are both iterable, treat them specially
            if isinstance(item, obj.NoneObject) or isinstance(item, str):
                return item

            if isinstance(item, types.GeneratorType):
                raise CacheContainsGenerator
            for x in iter(item):
                flat_x = self._find_generators(x)
                result.append(flat_x)

            return result
        except TypeError:
            return item

    def set_payload(self, payload):
        ''' Update the current payload with the new specified payload '''
        try:
            self.payload = self._find_generators(payload)
        except CacheContainsGenerator:
            # This only works because None payload cached results are rerun
            self.payload = None

    def dump(self):
        ''' Dump the node to disk for later retrieval. This is
        normally called when the process has exited. '''
        if self.payload:
            self.storage.dump(self.stem, self)

    def get_payload(self):
        """Retrieve this node's payload"""
        return self.payload

class BlockingNode(CacheNode):
    """Node that fails on all cache attempts and no-ops on cache storage attempts"""
    def __init__(self, name, stem, **kwargs):
        CacheNode.__init__(self, name, stem, **kwargs)

    def __getitem__(self, item = ''):
        return BlockingNode(item, '/'.join((self.stem, item)))

    def dump(self):
        """Ensure nothing gets dumped"""
        pass

    def get_payload(self):
        """Do not set a payload for a blocked cache node"""
        pass

class Invalidator(object):
    """ The Invalidator encapsulates program state to control
    invalidation of the cache.

    1) This object registers callbacks using the add_condition()
    method.

    2) Prior to serialising the cache object the callbacks are called
    returning a signature dict.

    3) When unpickling the cached object, we call the invalidator to
    produce a signature dict again, and compare this to the pickled
    version.

    The purpose of the callbacks is to represent a signature of the
    current state of execution. If the signature changes, the cache is
    invalidated.
    """
    def __init__(self):
        self.callbacks = {}

    def add_condition(self, key, callback):
        """Callback will be stored under key and should return a string.
        """
        self.callbacks[key] = callback

    def __setstate__(self, state):
        ## We do not actually have any callbacks here - we must use
        ## the global cache invalidator. We cant really get away from
        ## having a global invalidator.
        for k, v in CACHE.invalidator.callbacks.items():
            # TODO: Determine what happens if the state or current callbacks
            # contain a key that's not in the other
            if k in state and v() != state[k]:
                debug.debug("Invaliding cache... {0} (Running) != {1} (Stored) on key {2}".format(v(), state[k], k))

                raise InvalidCache("Running environment inconsistant "
                                   "with pickled environment - "
                                   "invalidating cache.")

    def __getstate__(self):
        """When pickling ourselves we call our callbacks to provide a
        dict of strings (our state signature). This dict should
        reflect all of our running state at the moment. This will then
        be compared to the state signature when unpickling and if its
        different we invalidate the cache.
        """
        result = {}
        for k, v in CACHE.invalidator.callbacks.items():
            result[k] = v()

        debug.debug("Pickling State signature: {0}".format(result))

        return result

class CacheTree(object):
    """ An abstract structure which represents the cache tree """
    def __init__(self, storage = None, cls = CacheNode, invalidator = None):
        self.storage = storage
        self.cls = cls
        self.invalidator = invalidator
        self.root = self.cls('', '', storage = storage, invalidator = invalidator)

    def __getitem__(self, path):
        """Pythonic interface to the cache"""
        return self.check(path, cls = self.cls)

    def invalidate_on(self, key, callback):
        self.invalidator.add_condition(key, callback)

    def check(self, path, callback = None, cls = CacheNode):
        """ Retrieves the node at the path specified """
        # Abort if we haven't been given a location
        if not config.LOCATION:
            return None

        ## Normalise the path
        path = urlparse.urljoin(config.LOCATION + "/", path)

        elements = path.split("/")
        current = self.root

        for e in elements:
            try:
                current = current[e]
            except KeyError:
                if current.stem:
                    next_stem = '/'.join((current.stem, e))
                else:
                    next_stem = e

                payload = None
                if callback is not None:
                    payload = callback()

                node = cls(e, next_stem, storage = self.storage,
                           payload = payload, invalidator = self.invalidator)

                current = node

        return current

class CacheStorage(object):
    """ The base class for implementation storing the cache. """
    ## Characters allowed in filenames (/'s are allowed since we're dealing with URLs only)
    printables = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-_./"

    def encode(self, string):
        result = ''
        for x in string:
            if x in self.printables:
                result += x
            else:
                result += "%{0:02X}".format(ord(x))

        return result

    def filename(self, url):
        if url.startswith(config.LOCATION):
            # Encode just the path part, since everything else is taken from relatively safe/already used data
            path = self.encode(url[len(config.LOCATION):])
        else:
            raise exceptions.CacheRelativeURLException("Storing non relative URLs is not supported now ({0})".format(url))

        # Join together the bits we need, and abspath it to ensure it's right for the OS it's on
        path = os.path.abspath(os.path.sep.join([config.CACHE_DIRECTORY,
                                                 os.path.basename(config.LOCATION) + ".cache",
                                                 path + '.pickle']))

        return path

    def load(self, url):
        filename = self.filename(url)

        debug.debug("Loading from {0}".format(filename))
        data = open(filename).read()

        debug.trace(level = 3)
        return pickle.loads(data)

    def dump(self, url, payload):
        # TODO: Ensure a better check for ieee1394/non-cachable address spaces than a bad URL
        try:
            filename = self.filename(url)
        except exceptions.CacheRelativeURLException:
            debug.debug("NOT Dumping url {0} - relative URLs are not yet supported".format(url))
            return

        ## Check that the directory exists
        directory = os.path.dirname(filename)
        if not os.access(directory, os.R_OK | os.W_OK | os.X_OK):
            os.makedirs(directory)

        ## Ensure that the payload is flattened - i.e. all generators are converted to lists for pickling
        try:
            data = pickle.dumps(payload)
            debug.debug("Dumping filename {0}".format(filename))
            fd = open(filename, 'w')
            fd.write(data)
            fd.close()
        except (pickle.PickleError, TypeError):
            # Do nothing if the pickle fails
            debug.debug("NOT Dumping filename {0} - contained a non-picklable class".format(filename))

## This is the central cache object
CACHE = CacheTree(CacheStorage(), BlockingNode, invalidator = Invalidator())

def enable_caching(_option, _opt_str, _value, _parser):
    """Turns off caching by replacing the tree with one that only takes BlockingNodes"""
    debug.debug("Enabling Caching")
    # Feels filthy using the global keyword,
    # but I can't figure another way to ensure that
    # the code gets called and overwrites the outer scope
    global CACHE
    CACHE = CacheTree(CacheStorage(), invalidator = Invalidator())
    config.CACHE = True

config.add_option("CACHE", default = False, action = 'callback',
                  cache_invalidator = False,
                  callback = enable_caching,
                  help = "Use caching")

class CacheDecorator(object):
    """ This decorator will memoise a function in the cache """
    def __init__(self, path):
        """Wraps a function in a cache decorator.

        The results of the function will be cached and memoised. Further
        calls to the function will retrieve the result from the
        cache. Cached objects are stored with the specified path as a
        key.

        Args:
           path: Key for storage into the cache. If this is callable,
           it will be called with the function's args and is expected
           to return a string which will be used as a path.

        Returns:
           A decorator.

        Example: Suppose the calculate function is decorated:

        @CacheDecorator(lambda self: "tests/pslist/pid{0}/".format(self._config.PID))
        def calculate(self):
           ....

        Note the use of the callback to finely tune the cache key depending on external variables.
        """
        self.path = path
        self.node = None

    def generate(self, path, g):
        """ Special handling for generators. We pass each iteration
        back immediately, and keep it in a list. Note that if the
        generator is aborted, the cache is not dumped.
        """
        payload = []
        for x in g:
            payload.append(x)
            yield x

        self.dump(path, payload)

    def dump(self, path, payload):
        self.node = CACHE[path]
        self.node.set_payload(payload)
        self.node.dump()

    def _cachewrapper(self, f, s, *args, **kwargs):
        """Wrapper for caching function calls"""
        ## See if the path is callable:
        if callable(self.path):
            path = self.path(s, *args, **kwargs)
        else:
            path = self.path

        ## Check if the result can be retrieved
        self.node = CACHE[path]
        # If this test goes away, we need to change the set_payload exception check
        # to act on dump instead of just the payload
        if self.node:
            payload = self.node.get_payload()
            if payload:
                return payload

        result = f(s, *args, **kwargs)

        ## If the wrapped function is a generator we need to
        ## handle it especially
        if isinstance(result, types.GeneratorType):
            return self.generate(path, result)

        self.dump(path, result)
        return result

    def __call__(self, f):
        def wrapper(s, *args, **kwargs):
            if config.CACHE:
                return self._cachewrapper(f, s, *args, **kwargs)

            return f(s, *args, **kwargs)

        return wrapper

class TestDecorator(CacheDecorator):
    """This decorator is just like a CacheDecorator, but will *always* cache fully"""

    def __call__(self, f):
        def wrapper(s, *args, **kwargs):
            return self._cachewrapper(f, s, *args, **kwargs)
        return wrapper

class Testable(object):
    """ This is a mixin that makes a class response to the unit tests 
    
        It must be inheritted *after* the command class
    """

    def calculate(self):
        """Empty function used to allow mixin"""

    def _flatten(self, item):
        """Flattens an item, including all generators"""
        try:
            # Make sure dicts aren't flattened to lists
            if isinstance(item, dict):
                result = {}
                for i in item:
                    result[self._flatten(i)] = self._flatten(item[i])
                return result

            for x in iter(item):
                flat_x = self._flatten(x)

            return flat_x
        except TypeError:
            return item

    ## This forces the test to be memoised with a key name derived from the class name
    @TestDecorator(lambda self: "tests/unittests/{0}".format(self.__class__.__name__))
    def test(self):
        ## This forces iteration over all keys - this is required in order
        ## to flatten the full list for the cache
        ## We must ensure config.CACHE is False here, otherwise the change isn't registered in this module
        config.CACHE = False
        return self._flatten(self.calculate())