Repository: google/turbinia
Pull request #476 (2019-12-07 00:20:30)

Add graph engine option to graph generation
Also make magic "list" value that will list out the engine and output format types.


Patch:
diff --git a/tools/turbinia_job_graph.py b/tools/turbinia_job_graph.py
index 72a404c..dc577b9 100755
--- a/tools/turbinia_job_graph.py
+++ b/tools/turbinia_job_graph.py
@@ -1,4 +1,4 @@
-#!/usr/bin/env python -v
+#!/usr/bin/env python
 # -*- coding: utf-8 -*-
 # Copyright 2018 Google Inc.
 #
@@ -54,16 +54,36 @@ if __name__ == '__main__':
       description='Create Turbinia evidence graph.')
   parser.add_argument(
       '-f', '--format', default='png',
-      help='The format of the output file you wish to generate.  Supported '
-      'types are here: http://www.graphviz.org/doc/info/output.html')
+      help='The format of the output file you wish to generate.  Specify '
+      '"list" to list out the available output types.  More info is here: '
+      'http://www.graphviz.org/doc/info/output.html')
+  parser.add_argument(
+      '-e', '--engine', default='dot',
+      help='The graphviz engine used to generate the graph layout.  Specify '
+      '"list" to list out the available engines.')
   parser.add_argument('filename', type=unicode, help='where to save the file')
   args = parser.parse_args()
 
+  if args.format == 'list':
+    formats = ' '.join(graphviz.FORMATS)
+    print('Available format types: {0:s}'.format(formats))
+    sys.exit(0)
+
   if args.format not in graphviz.FORMATS:
     print('Format type {0:s} is not supported'.format(args.format))
     sys.exit(1)
 
+  if args.engine == 'list':
+    engines = ' '.join(graphviz.ENGINES)
+    print('Available graph layout engines: {0:s}'.format(engines))
+    sys.exit(0)
+
+  if args.engine not in graphviz.ENGINES:
+    print('Layout engine type {0:s} is not supported'.format(args.engine))
+    sys.exit(1)
+
   graph = create_graph()
+  graph.engine = args.engine
   output_file = args.filename.replace('.png', '')
 
   try:
diff --git a/turbinia/config/__init__.py b/turbinia/config/__init__.py
index 93f9792..8bc2745 100644
--- a/turbinia/config/__init__.py
+++ b/turbinia/config/__init__.py
@@ -163,13 +163,24 @@ def ValidateAndSetConfig(_config):
 
   CONFIGVARS = REQUIRED_VARS + OPTIONAL_VARS
   for var in CONFIGVARS:
+    empty_value = False
     if not hasattr(_config, var):
-      raise TurbiniaException(
-          'No config attribute {0:s}:{1:s}'.format(_config.configSource, var))
+      if var in OPTIONAL_VARS:
+        log.debug(
+            'Setting non-existent but optional config variable {0:s} to '
+            'None'.format(var))
+        empty_value = True
+      else:
+        raise TurbiniaException(
+            'Required config attribute {0:s}:{1:s} not in config'.format(
+                _config.configSource, var))
     if var in REQUIRED_VARS and getattr(_config, var) is None:
       raise TurbiniaException(
           'Config attribute {0:s}:{1:s} is not set'.format(
               _config.configSource, var))
 
     # Set the attribute in the current module
-    setattr(sys.modules[__name__], var, getattr(_config, var))
+    if empty_value:
+      setattr(sys.modules[__name__], var, None)
+    else:
+      setattr(sys.modules[__name__], var, getattr(_config, var))
diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
index 9516f62..657d28c 100644
--- a/turbinia/workers/__init__.py
+++ b/turbinia/workers/__init__.py
@@ -613,7 +613,7 @@ class TurbiniaTask(object):
                   self.turbinia_version, turbinia.__version__))
           self.result.log(message, level=logging.ERROR)
           self.result.status = message
-          return self.result
+          return self.result.serialize()
 
         self._evidence_config = evidence.config
         self.result = self.run(evidence, self.result)


Test patch:
diff --git a/turbinia/config/turbinia_config_test.py b/turbinia/config/turbinia_config_test.py
index 2a63456..51e1079 100644
--- a/turbinia/config/turbinia_config_test.py
+++ b/turbinia/config/turbinia_config_test.py
@@ -110,6 +110,15 @@ class TestTurbiniaConfig(unittest.TestCase):
     self.assertEqual(config.UNSETKEY, None)
     self.assertEqual(config.SETKEY, 'bar')
 
+  def testNonexistentOptionalKeyConfig(self):
+    """Test that optional vars don't need to exist."""
+    self.WriteConfig('SETKEY = "bar"\n')
+    config.REQUIRED_VARS = ['SETKEY']
+    config.OPTIONAL_VARS = ['UNSETKEY']
+    config.LoadConfig()
+    self.assertEqual(config.UNSETKEY, None)
+    self.assertEqual(config.SETKEY, 'bar')
+
   def testMissingConfigPath(self):
     """Test non-existent config path."""
     config.CONFIGPATH = ['DOESNOTEXIST']