From c605cbd18254e153d31aa88eba8acfb9bc0a80c8 Mon Sep 17 00:00:00 2001
From: Mike Auty <mike.auty@gmail.com>
Date: Thu, 8 Mar 2018 01:20:34 +0000
Subject: [PATCH] Add in TZAware and TZNaive values for TreeGrids.

---
 volatility/framework/interfaces/renderers.py | 16 +++++++++++++++-
 volatility/framework/renderers/__init__.py   | 17 +++++++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/volatility/framework/interfaces/renderers.py b/volatility/framework/interfaces/renderers.py
index cb9ec1407..533f6eac0 100644
--- a/volatility/framework/interfaces/renderers.py
+++ b/volatility/framework/interfaces/renderers.py
@@ -3,6 +3,7 @@ or in some other form.  This module defines both the output format (:class:`Tree
 which can interact with a TreeGrid to produce suitable output."""
 
 import collections
+import datetime
 import typing
 from abc import abstractmethod, ABCMeta
 
@@ -77,12 +78,25 @@ class BaseAbsentValue(object):
     """Class that represents values which are not present for some reason"""
 
 
+# We don't class these off a shared base, because the SimpleTypes must only
+# contain the types that the validator will accept (which would not include the base)
+
+class TZAwareValue(datetime.datetime):
+    """Class for TZ-aware datetimes"""
+
+
+class TZNaiveValue(datetime.datetime):
+    """Class for TZ-aware datetimes"""
+
+
 _Type = typing.TypeVar("_Type")
 ColumnsType = typing.List[typing.Tuple[str, typing.Type]]
 SimpleTypes = typing.Union[typing.Type[int],
                            typing.Type[str],
                            typing.Type[float],
                            typing.Type[bytes],
+                           typing.Type[TZAwareValue],
+                           typing.Type[TZNaiveValue],
                            typing.Type[BaseAbsentValue]]
 VisitorSignature = typing.Callable[[TreeNode, _Type], _Type]
 
@@ -100,7 +114,7 @@ class TreeGrid(object, metaclass = ABCMeta):
     and to create cycles.
     """
 
-    simple_types = (int, str, float, bytes)  # type: typing.ClassVar[typing.Tuple]
+    simple_types = (int, str, float, bytes, datetime)  # type: typing.ClassVar[typing.Tuple]
 
     def __init__(self, columns: ColumnsType, generator: typing.Generator) -> None:
         """Constructs a TreeGrid object using a specific set of columns
diff --git a/volatility/framework/renderers/__init__.py b/volatility/framework/renderers/__init__.py
index aabb141d7..27fddc75a 100644
--- a/volatility/framework/renderers/__init__.py
+++ b/volatility/framework/renderers/__init__.py
@@ -52,13 +52,26 @@ class TreeNode(interfaces.renderers.TreeNode):
                 "Values must be a list of objects made up of simple types and number the same as the columns")
         for index in range(len(self._treegrid.columns)):
             column = self._treegrid.columns[index]
-            if not isinstance(values[index], (column.type, interfaces.renderers.BaseAbsentValue)):
+            val = values[index]
+            if not isinstance(val, (column.type, interfaces.renderers.BaseAbsentValue)):
                 raise TypeError(
                     "Values item with index {} is the wrong type for column {} (got {} but expected {})".format(
                         index,
                         column.name,
-                        type(values[index]),
+                        type(val),
                         column.type))
+            if isinstance(val, (interfaces.renderers.TZAwareValue, interfaces.renderers.TZNaiveValue)):
+                tznaive = val.tzinfo is None or val.tzinfo.utcoffset(val) is None
+                if isinstance(val, interfaces.renderers.TZAwareValue) and tznaive:
+                    raise TypeError(
+                        "Values item with index {} is not a timezone aware datetime object as required by column {}".format(
+                            index,
+                            column.name))
+                elif isinstance(val, interfaces.renderers.TZNaiveValue) and not tznaive:
+                    raise TypeError(
+                        "Values item with index {} is not a timezone naive datetime object as required by column {}".format(
+                            index,
+                            column.name))
 
     @property
     def values(self) -> typing.Iterable[interfaces.renderers.SimpleTypes]: