""" Created on 7 May 2013 @author: mike """ import collections import copy import json import logging from volatility.framework.configuration.requirements import MultiRequirement from volatility.framework.interfaces.configuration import CONFIG_SEPARATOR vollog = logging.getLogger(__name__) class HierarchicalDict(collections.Mapping): def __init__(self, initial_dict = None, separator = CONFIG_SEPARATOR): if not (isinstance(separator, str) and len(separator) == 1): raise TypeError("Separator must be a one character string") self._separator = separator self._data = {} self._subdict = {} if isinstance(initial_dict, str): initial_dict = json.loads(initial_dict) if isinstance(initial_dict, dict): for k, v in initial_dict.items(): self[k] = v elif initial_dict is not None: raise TypeError("Initial_dict must be a dictionary or JSON string containing a dictionary") @property def separator(self): return self._separator @property def data(self): return self._data.copy() def _key_head(self, key): """Returns the first division of a key based on the dict separator, or the full key if the separator is not present """ if self.separator in key: return key[:key.index(self.separator)] else: return key def _key_tail(self, key): """Returns all but the first division of a key based on the dict separator, or None if the separator is not in the key """ if self.separator in key: return key[key.index(self.separator) + 1:] return None def __iter__(self): """Returns an iterator object that supports the iterator protocol""" return self.generator() def generator(self): """Yields the next element in the iterator""" for key in self._data: yield key for subdict_key in self._subdict: for key in self._subdict[subdict_key]: yield subdict_key + self.separator + key def __getitem__(self, key): """Gets an item, traversing down the trees to get to the final value""" try: if self.separator in key: subdict = self._subdict[self._key_head(key)] return subdict[self._key_tail(key)] else: return self._data[key] except KeyError: raise KeyError(key) def __setitem__(self, key, value): """Sets an item or creates a subdict and sets the item within that""" self._setitem(key, value) def _setitem(self, key, value, is_data = True): """Set an item or appends a whole subtree at a key location""" if self.separator in key: subdict = self._subdict.get(self._key_head(key), HierarchicalDict(separator = self.separator)) subdict[self._key_tail(key)] = value self._subdict[self._key_head(key)] = subdict else: if is_data: self._data[key] = value else: if not isinstance(value, HierarchicalDict) and value is not None: raise TypeError("HierarchicalDicts can only store HierarchicalDicts within their structure") self._subdict[key] = value def __delitem__(self, key): """Deletes an item from the hierarchical dict""" try: if self.separator in key: subdict = self._subdict[self._key_head(key)] del subdict[self._key_tail(key)] if not subdict: del self._subdict[self._key_head(key)] except KeyError: raise KeyError(key) def __contains__(self, key): """Determines whether the key is present in the hierarchy""" if self.separator in key: try: subdict = self._subdict[self._key_head(key)] return self._key_tail(key) in subdict except KeyError: return False else: return key in self._data def __len__(self): """Returns the length of all items""" return len(self._data) + sum([len(subdict) for subdict in self._subdict]) def branch(self, key): """Returns the HierarchicalDict housed under the key""" if self.separator in key: return self._subdict[self._key_head(key)].branch(self._key_tail(key)) else: return self._subdict[key] def splice(self, key, subdict): """Overwrites a branch in an existing tree with a new tree WARNING: This will *NOT* make alterations to objects that have already been given a config_path meaning those objects will fail their configuration lookup. This can also be used to bulk delete a branch of a tree by passing in None """ self._setitem(key, subdict, is_data = False) def clone(self): """Duplicate the configuration, allowing changes without affecting the original""" return copy.deepcopy(self) def __str__(self): """Turns the Hierarchical dict into a string representation""" return json.dumps(dict([(key, self[key]) for key in self.generator()]), indent = 2)