Repository: google/turbinia
Pull request #464 (2019-11-27 17:24:19)

Added functionality to filter jobs and run a system dependency check on the worker. 
Added functionality to filter jobs by worker via --jobs_whitelist or --jobs_blacklist flags or in the config file inside the DISABLE_JOBS variable. Updated the config to do a system dependency check prior to starting the worker (it will check if the job name in the dependency check matches the list of available jobs and then run a check for all the programs to see if they exist or not.). The task.run_wrapper will check if a given job is enabled in the worker before running and if it does not, will set the result.status warning the user that the job has been disabled on the worker. 


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/client.py b/turbinia/client.py
index 40d5975..e5a3796 100644
--- a/turbinia/client.py
+++ b/turbinia/client.py
@@ -19,6 +19,7 @@ from __future__ import unicode_literals
 from collections import defaultdict
 from datetime import datetime
 from datetime import timedelta
+
 import json
 import logging
 from operator import itemgetter
@@ -26,6 +27,7 @@ from operator import attrgetter
 import os
 import stat
 import time
+import shutil
 
 from turbinia import config
 from turbinia.config import logger
@@ -33,6 +35,7 @@ from turbinia.config import DATETIME_FORMAT
 from turbinia import task_manager
 from turbinia import TurbiniaException
 from turbinia.lib import text_formatter as fmt
+from turbinia.jobs import manager as job_manager
 from turbinia.workers import Priority
 from turbinia.workers.artifact import FileArtifactExtractionTask
 from turbinia.workers.analysis.wordpress import WordpressAccessLogAnalysisTask
@@ -90,6 +93,43 @@ log = logging.getLogger('turbinia')
 logger.setup()
 
 
+def check_dependencies(dependencies):
+  """Checks system dependencies.
+
+  Args:
+    dependencies(dict): dictionary of dependencies to check for.
+
+  Raises:
+    TurbiniaException: if dependency is not met or a bad config file.
+  """
+
+  log.info('Performing system dependency check.')
+  job_names = list(job_manager.JobsManager.GetJobNames())
+  # Iterate through dependency config and perform a system check.
+  try:
+    for dep in dependencies:
+      if dep['job'].lower() in job_names:
+        for p in dep['programs']:
+          # Program can not be found.
+          if shutil.which(p) is None:
+            raise TurbiniaException(
+                'Job dependency {0:s} not found for job: {1:s}. Please install '
+                'the dependency or disable the job.'.format(p, dep['job']))
+      # If job is not found in list of registered jobs.
+      else:
+        log.warning(
+            'The job: {0:s} was not found or has been disabled. Skipping '
+            'dependency check...'.format(dep['job']))
+  except (KeyError, TypeError) as exception:
+    raise TurbiniaException(
+        'An issue has occured while parsing the '
+        'dependency config:{0!s} '.format(exception))
+
+  log.info(
+      'Dependency check complete. The following jobs will be active:'
+      'for this worker: {0:s}'.format(', '.join(job_names)))
+
+
 def check_directory(directory):
   """Checks directory to make sure it exists and is writable.
 
@@ -230,7 +270,7 @@ class TurbiniaClient(object):
     # the task manager from the client.
     log.info('Available Jobs:')
     for job in self.task_manager.jobs:
-      log.info('\t{0:s}'.format(job.name))
+      log.info('\t{0:s}'.format(job.NAME))
 
   def wait_for_request(
       self, instance, project, region, request_id=None, user=None,
@@ -783,9 +823,25 @@ class TurbiniaCeleryWorker(TurbiniaClient):
     worker (celery.app): Celery worker app
   """
 
-  def __init__(self, *_, **__):
-    """Initialization for Celery worker."""
+  def __init__(self, jobs_blacklist=None, jobs_whitelist=None):
+    """Initialization for celery worker.
+
+    Args:
+      jobs_blacklist (Optional[list[str]]): Jobs we will exclude from running
+      jobs_whitelist (Optional[list[str]]): The only Jobs we will include to run
+    """
     super(TurbiniaCeleryWorker, self).__init__()
+    # Deregister jobs from blacklist/whitelist.
+    disabled_jobs = list(config.DISABLED_JOBS) if config.DISABLED_JOBS else []
+    job_manager.JobsManager.DeregisterJobs(jobs_blacklist, jobs_whitelist)
+    if disabled_jobs:
+      log.info(
+          'Disabling jobs that were configured to be disabled in the '
+          'config file: {0:s}'.format(', '.join(disabled_jobs)))
+      job_manager.JobsManager.DeregisterJobs(jobs_blacklist=disabled_jobs)
+
+    # Check for valid dependencies/directories.
+    check_dependencies(config.DEPENDENCIES)
     check_directory(config.MOUNT_DIR_PREFIX)
     check_directory(config.OUTPUT_DIR)
     check_directory(config.TMP_DIR)
@@ -810,8 +866,13 @@ class TurbiniaPsqWorker(object):
     TurbiniaException: When errors occur
   """
 
-  def __init__(self, *_, **__):
-    """Initialization for PSQ Worker."""
+  def __init__(self, jobs_blacklist=None, jobs_whitelist=None):
+    """Initialization for PSQ Worker.
+
+    Args:
+      jobs_blacklist (Optional[list[str]]): Jobs we will exclude from running
+      jobs_whitelist (Optional[list[str]]): The only Jobs we will include to run
+    """
     config.LoadConfig()
     psq_publisher = pubsub.PublisherClient()
     psq_subscriber = pubsub.SubscriberClient()
@@ -825,6 +886,17 @@ class TurbiniaPsqWorker(object):
       log.error(msg)
       raise TurbiniaException(msg)
 
+    # Deregister jobs from blacklist/whitelist.
+    disabled_jobs = list(config.DISABLED_JOBS) if config.DISABLED_JOBS else []
+    job_manager.JobsManager.DeregisterJobs(jobs_blacklist, jobs_whitelist)
+    if disabled_jobs:
+      log.info(
+          'Disabling jobs that were configured to be disabled in the '
+          'config file: {0:s}'.format(', '.join(disabled_jobs)))
+      job_manager.JobsManager.DeregisterJobs(jobs_blacklist=disabled_jobs)
+
+    # Check for valid dependencies/directories.
+    check_dependencies(config.DEPENDENCIES)
     check_directory(config.MOUNT_DIR_PREFIX)
     check_directory(config.OUTPUT_DIR)
     check_directory(config.TMP_DIR)
diff --git a/turbinia/config/__init__.py b/turbinia/config/__init__.py
index 93f9792..acac734 100644
--- a/turbinia/config/__init__.py
+++ b/turbinia/config/__init__.py
@@ -55,6 +55,9 @@ REQUIRED_VARS = [
     'SHARED_FILESYSTEM',
     # TODO(aarontp): Move this to the recipe config when it's available.
     'DEBUG_TASKS',
+    'DEPENDENCIES',
+    'DOCKER_ENABLED',
+    'DISABLED_JOBS',
 ]
 
 # Optional config vars.  Some may be mandatory depending on the configuration
@@ -163,13 +166,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/config/turbinia_config_tmpl.py b/turbinia/config/turbinia_config_tmpl.py
index 4bb004e..b431f7d 100644
--- a/turbinia/config/turbinia_config_tmpl.py
+++ b/turbinia/config/turbinia_config_tmpl.py
@@ -73,6 +73,70 @@ SHARED_FILESYSTEM = False
 # problems.
 DEBUG_TASKS = False
 
+################################################################################
+#                         External Dependency Configurations
+#
+# Options in this section are used to configure system and docker dependencies.
+################################################################################
+
+# This will allow for the configuration of system dependencies and docker
+# containers. The following configuration will need to be set per job config.
+# Please use the example below as a guide to add additional job checks.
+# {
+#   'job': 'MyJob'
+#   'programs': ['list', 'of', 'required', 'programs']
+#   'docker_image': 'ImageID' # Or None if no image configured for this job.
+# }
+
+# This will enable the usage of docker containers for the worker.
+DOCKER_ENABLED = False
+
+# Any jobs added to this list will disable it from being used.
+DISABLED_JOBS = []
+
+# Configure additional job dependency checks below.
+DEPENDENCIES = [{
+    'job': 'BinaryExtractorJob',
+    'programs': ['image_export.py'],
+    'docker_image': None
+}, {
+    'job': 'BulkExtractorJob',
+    'programs': ['bulk_extractor'],
+    'docker_image': None
+}, {
+    'job': 'GrepJob',
+    'programs': ['grep'],
+    'docker_image': None
+}, {
+    'job': 'HadoopAnalysisJob',
+    'programs': ['strings'],
+    'docker_image': None
+}, {
+    'job': 'HindsightJob',
+    'programs': ['hindsight.py'],
+    'docker_image': None
+}, {
+    'job': 'JenkinsAnalysisJob',
+    'programs': ['john'],
+    'docker_image': None
+}, {
+    'job': 'PlasoJob',
+    'programs': ['log2timeline.py'],
+    'docker_image': None
+}, {
+    'job': 'PsortJob',
+    'programs': ['psort.py'],
+    'docker_image': None
+}, {
+    'job': 'StringsJob',
+    'programs': ['strings'],
+    'docker_image': None
+}, {
+    'job': 'VolatilityJob',
+    'programs': ['vol.py'],
+    'docker_image': None
+}]
+
 ################################################################################
 #                        Google Cloud Platform (GCP)
 #
diff --git a/turbinia/jobs/manager.py b/turbinia/jobs/manager.py
index 53fd6b1..1f86f63 100644
--- a/turbinia/jobs/manager.py
+++ b/turbinia/jobs/manager.py
@@ -93,6 +93,39 @@ class JobsManager(object):
 
     del cls._job_classes[job_name]
 
+  @classmethod
+  def DeregisterJobs(cls, jobs_blacklist=None, jobs_whitelist=None):
+    """Deregisters a list of job names against white/black lists.
+
+    jobs_whitelist and jobs_blacklist must not be specified at the same time.
+
+    Args:
+      jobs_blacklist (Optional[list[str]]): Job names to deregister.
+      jobs_whitelist (Optional[list[str]]): Job names to register.
+
+    Raises:
+      TurbiniaException if both jobs_blacklist and jobs_whitelist are specified.
+    """
+    registered_jobs = list(cls.GetJobNames())
+    jobs_remove = []
+
+    # Create a list of jobs to deregister.
+    if jobs_whitelist and jobs_blacklist:
+      raise TurbiniaException(
+          'jobs_whitelist and jobs_blacklist cannot be specified at the same '
+          'time.')
+    elif jobs_whitelist:
+      jobs_whitelist = [j.lower() for j in jobs_whitelist]
+      jobs_remove = [j for j in registered_jobs if j not in jobs_whitelist]
+    elif jobs_blacklist:
+      jobs_blacklist = [j.lower() for j in jobs_blacklist]
+      jobs_remove = [j for j in jobs_blacklist if j in registered_jobs]
+
+    # Deregister the jobs.
+    jobs_remove = [j.lower() for j in jobs_remove]
+    for job_name in jobs_remove:
+      del cls._job_classes[job_name]
+
   @classmethod
   def GetJobInstance(cls, job_name):
     """Retrieves an instance of a specific job.
diff --git a/turbinia/task_manager.py b/turbinia/task_manager.py
index 85fc96e..24d65d3 100644
--- a/turbinia/task_manager.py
+++ b/turbinia/task_manager.py
@@ -313,6 +313,7 @@ class BaseTaskManager(object):
     task.requester = evidence_.config.get('requester')
     if job:
       task.job_id = job.id
+      task.job_name = job.name
       job.tasks.append(task)
     self.state_manager.write_new_task(task)
     self.enqueue_task(task, evidence_)
diff --git a/turbinia/turbiniactl.py b/turbinia/turbiniactl.py
index 6d6b404..29faaa6 100644
--- a/turbinia/turbiniactl.py
+++ b/turbinia/turbiniactl.py
@@ -576,11 +576,13 @@ def main():
     # Set up root logger level which is normally set by the psqworker command
     # which we are bypassing.
     logger.setup()
-    worker = TurbiniaPsqWorker()
+    worker = TurbiniaPsqWorker(
+        jobs_blacklist=args.jobs_blacklist, jobs_whitelist=args.jobs_whitelist)
     worker.start()
   elif args.command == 'celeryworker':
     logger.setup()
-    worker = TurbiniaCeleryWorker()
+    worker = TurbiniaCeleryWorker(
+        jobs_blacklist=args.jobs_blacklist, jobs_whitelist=args.jobs_whitelist)
     worker.start()
   elif args.command == 'server':
     server = TurbiniaServer(
diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
index 9516f62..d5ea9db 100644
--- a/turbinia/workers/__init__.py
+++ b/turbinia/workers/__init__.py
@@ -306,6 +306,7 @@ class TurbiniaTask(object):
       id (str): Unique Id of task (string of hex)
       is_finalize_task (bool): Whether this is a finalize Task or not.
       job_id (str): Job ID the Task was created by.
+      job_name (str): The name of the Job.
       last_update (datetime): A datetime object with the last time the task was
           updated.
       name (str): Name of task
@@ -342,6 +343,7 @@ class TurbiniaTask(object):
     self.id = uuid.uuid4().hex
     self.is_finalize_task = False
     self.job_id = None
+    self.job_name = None
     self.last_update = datetime.now()
     self.name = name if name else self.__class__.__name__
     self.output_dir = None
@@ -597,6 +599,9 @@ class TurbiniaTask(object):
     Returns:
       A TurbiniaTaskResult object
     """
+    # Avoid circular dependency.
+    from turbinia.jobs import manager as job_manager
+
     log.debug('Task {0:s} {1:s} awaiting execution'.format(self.name, self.id))
     evidence = evidence_decode(evidence)
     with filelock.FileLock(config.LOCK_FILE):
@@ -607,13 +612,26 @@ class TurbiniaTask(object):
         original_result_id = self.result.id
         evidence.validate()
 
+        # TODO(wyassine): refactor it so the result task does not
+        # have to go through the preprocess stage. At the moment
+        # self.results.setup is required to be called to set its status.
+        # Check if Task's job is available for the worker.
+        active_jobs = list(job_manager.JobsManager.GetJobNames())
+        if self.job_name.lower() not in active_jobs:
+          message = (
+              'Task will not run due to the job: {0:s} being disabled '
+              'on the worker.'.format(self.job_name))
+          self.result.log(message, level=logging.ERROR)
+          self.result.status = message
+          return self.result.serialize()
+
         if self.turbinia_version != turbinia.__version__:
           message = (
               'Worker and Server versions do not match: {0:s} != {1:s}'.format(
                   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/client_test.py b/turbinia/client_test.py
index 24c9a64..b27626f 100644
--- a/turbinia/client_test.py
+++ b/turbinia/client_test.py
@@ -32,6 +32,7 @@ from turbinia.client import TurbiniaClient
 from turbinia.client import TurbiniaServer
 from turbinia.client import TurbiniaStats
 from turbinia.client import TurbiniaPsqWorker
+from turbinia.client import check_dependencies
 from turbinia import TurbiniaException
 
 SHORT_REPORT = textwrap.dedent(
@@ -440,6 +441,7 @@ class TestTurbiniaPsqWorker(unittest.TestCase):
     config.LoadConfig()
     config.OUTPUT_DIR = self.tmp_dir
     config.MOUNT_DIR_PREFIX = self.tmp_dir
+    config.DEPENDENCIES = []
 
   def tearDown(self):
     if 'turbinia-test' in self.tmp_dir:
@@ -450,7 +452,7 @@ class TestTurbiniaPsqWorker(unittest.TestCase):
   @mock.patch('turbinia.client.psq.Worker')
   def testTurbiniaPsqWorkerInit(self, _, __, ___):
     """Basic test for PSQ worker."""
-    worker = TurbiniaPsqWorker()
+    worker = TurbiniaPsqWorker([], [])
     self.assertTrue(hasattr(worker, 'worker'))
 
   @mock.patch('turbinia.client.pubsub')
@@ -459,7 +461,7 @@ class TestTurbiniaPsqWorker(unittest.TestCase):
   def testTurbiniaClientNoDir(self, _, __, ___):
     """Test that OUTPUT_DIR path is created."""
     config.OUTPUT_DIR = os.path.join(self.tmp_dir, 'no_such_dir')
-    TurbiniaPsqWorker()
+    TurbiniaPsqWorker([], [])
     self.assertTrue(os.path.exists(config.OUTPUT_DIR))
 
   @mock.patch('turbinia.client.pubsub')
@@ -470,3 +472,31 @@ class TestTurbiniaPsqWorker(unittest.TestCase):
     config.OUTPUT_DIR = os.path.join(self.tmp_dir, 'empty_file')
     open(config.OUTPUT_DIR, 'a').close()
     self.assertRaises(TurbiniaException, TurbiniaPsqWorker)
+
+  @mock.patch('turbinia.client.shutil')
+  @mock.patch('logging.Logger.warning')
+  def testDependencyCheck(self, mock_logger, mock_shutil):
+    """Test system dependency check."""
+    dependencies = [{
+        'job': 'PlasoJob',
+        'programs': ['non_exist'],
+        'docker_image': None
+    }]
+
+    # Dependency not found.
+    mock_shutil.which.return_value = None
+    self.assertRaises(TurbiniaException, check_dependencies, dependencies)
+
+    # Normal run.
+    mock_shutil.which.return_value = True
+    check_dependencies(dependencies)
+
+    # Job not found.
+    dependencies[0]['job'] = 'non_exist'
+    check_dependencies(dependencies)
+    mock_logger.assert_called_with(
+        'The job: non_exist was not found or has been disabled. '
+        'Skipping dependency check...')
+
+    # Bad dependency config.
+    self.assertRaises(TurbiniaException, check_dependencies, [{'test': 'test'}])
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']
diff --git a/turbinia/workers/workers_test.py b/turbinia/workers/workers_test.py
index b3cb350..53551d6 100644
--- a/turbinia/workers/workers_test.py
+++ b/turbinia/workers/workers_test.py
@@ -55,6 +55,7 @@ class TestTurbiniaTaskBase(unittest.TestCase):
     self.plaso_task.output_manager.get_local_output_dirs.return_value = (
         None, None)
     self.task = self.task_class(base_output_dir=self.base_output_dir)
+    self.task.job_name = 'PlasoJob'
     self.task.output_manager = mock.MagicMock()
     self.task.output_manager.get_local_output_dirs.return_value = (None, None)
 
@@ -150,6 +151,17 @@ class TestTurbiniaTask(TestTurbiniaTaskBase):
     self.assertEqual(type(new_result), TurbiniaTaskResult)
     self.assertIn('CheckedResult', new_result.status)
 
+  def testTurbiniaTaskJobUnavailable(self):
+    """Test that the run wrapper can fail if the job doesn't exist."""
+    self.setResults()
+    self.task.job_name = 'non_exist'
+    canary_status = (
+        'Task will not run due to the job: '
+        'non_exist being disabled on the worker.')
+    new_result = self.task.run_wrapper(self.evidence.__dict__)
+    new_result = TurbiniaTaskResult.deserialize(new_result)
+    self.assertEqual(new_result.status, canary_status)
+
   def testTurbiniaTaskRunWrapperExceptionThrown(self):
     """Test that the run wrapper recovers from run throwing an exception."""
     self.setResults()