text: # -*- coding: utf-8 -*-
  # Copyright 2016 Google Inc.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Task manager for Turbinia."""
  
  from __future__ import unicode_literals, absolute_import
  
  import logging
  import time
  
  import turbinia
  from turbinia import workers
  from turbinia import evidence
  from turbinia import config
  from turbinia import state_manager
  from turbinia import TurbiniaException
  from turbinia.jobs import manager as jobs_manager
  
  config.LoadConfig()
  if config.TASK_MANAGER.lower() == 'psq':
    import psq
  
    from google.cloud import exceptions
    from google.cloud import datastore
    from google.cloud import pubsub
  
    from turbinia import pubsub as turbinia_pubsub
  elif config.TASK_MANAGER.lower() == 'celery':
    from celery import states as celery_states
  
    from turbinia import celery as turbinia_celery
  
  log = logging.getLogger('turbinia')
  
  PSQ_TASK_TIMEOUT_SECONDS = 604800
  PSQ_QUEUE_WAIT_SECONDS = 2
  
  
  def get_task_manager():
    """Return task manager object based on config.
  
    Returns
      Initialized TaskManager object.
  
    Raises:
      TurbiniaException: When an unknown task manager type is specified
    """
    config.LoadConfig()
    # pylint: disable=no-else-return
    if config.TASK_MANAGER.lower() == 'psq':
      return PSQTaskManager()
    elif config.TASK_MANAGER.lower() == 'celery':
      return CeleryTaskManager()
    else:
      msg = 'Task Manager type "{0:s}" not implemented'.format(
          config.TASK_MANAGER)
      raise turbinia.TurbiniaException(msg)
  
  
  def task_runner(obj, *args, **kwargs):
    """Wrapper function to run specified TurbiniaTask object.
  
    Args:
      obj: An instantiated TurbiniaTask object.
      *args: Any Args to pass to obj.
      **kwargs: Any keyword args to pass to obj.
  
    Returns:
      Output from TurbiniaTask (should be TurbiniaTaskResult).
    """
    obj = workers.TurbiniaTask.deserialize(obj)
    return obj.run_wrapper(*args, **kwargs)
  
  
  class BaseTaskManager(object):
    """Class to manage Turbinia Tasks.
  
    Handles incoming new Evidence messages, adds new Tasks to the queue and
    processes results from Tasks that have run.
  
    Attributes:
      jobs (list[TurbiniaJob]): Uninstantiated job classes.
      running_jobs (list[TurbiniaJob]): A list of jobs that are
          currently running.
      evidence (list): A list of evidence objects to process.
      state_manager (DatastoreStateManager|RedisStateManager): State manager
          object to handle syncing with storage.
      tasks (list[TurbiniaTask]): Running tasks.
    """
  
    def __init__(self):
      self.jobs = []
      self.running_jobs = []
      self.state_manager = state_manager.get_state_manager()
  
    @property
    def tasks(self):
      """A property that returns all outstanding Tasks.
  
      Returns:
        list[TurbiniaTask]: All outstanding Tasks.
      """
      return [task for job in self.running_jobs for task in job.tasks]
  
    def _backend_setup(self, *args, **kwargs):
      """Sets up backend dependencies.
  
      Raises:
        TurbiniaException: When encountering fatal errors setting up dependencies.
      """
      raise NotImplementedError
  
    def setup(self, jobs_denylist=None, jobs_allowlist=None, *args, **kwargs):
      """Does setup of Task manager and its dependencies.
  
      Args:
        jobs_denylist (list): Jobs that will be excluded from running
        jobs_allowlist (list): The only Jobs will be included to run
      """
      self._backend_setup(*args, **kwargs)
      job_names = jobs_manager.JobsManager.GetJobNames()
      if jobs_denylist or jobs_allowlist:
        selected_jobs = jobs_denylist or jobs_allowlist
        for job in selected_jobs:
          if job.lower() not in job_names:
            msg = (
                'Error creating server. Job {0!s} is not found in registered '
                'jobs {1!s}.'.format(job, job_names))
            log.error(msg)
            raise TurbiniaException(msg)
        log.info(
            'Filtering Jobs with allowlist {0!s} and denylist {1!s}'.format(
                jobs_allowlist, jobs_denylist))
        job_names = jobs_manager.JobsManager.FilterJobNames(
            job_names, jobs_denylist, jobs_allowlist)
  
      # Disable any jobs from the config that were not previously allowlisted.
      disabled_jobs = list(config.DISABLED_JOBS) if config.DISABLED_JOBS else []
      disabled_jobs = [j.lower() for j in disabled_jobs]
      if jobs_allowlist:
        disabled_jobs = list(set(disabled_jobs) - set(jobs_allowlist))
      if disabled_jobs:
        log.info(
            'Disabling non-allowlisted jobs configured to be disabled in the '
            'config file: {0:s}'.format(', '.join(disabled_jobs)))
        job_names = jobs_manager.JobsManager.FilterJobNames(
            job_names, disabled_jobs, [])
  
      self.jobs = [job for _, job in jobs_manager.JobsManager.GetJobs(job_names)]
      log.debug('Registered job list: {0:s}'.format(str(job_names)))
  
    def add_evidence(self, evidence_):
      """Adds new evidence and creates tasks to process it.
  
      This creates all tasks configured to process the given type of evidence.
  
      Args:
        evidence_: evidence object to add.
  
      Raises:
        TurbiniaException: When no Jobs are found.
      """
      if not self.jobs:
        raise turbinia.TurbiniaException(
            'Jobs must be registered before evidence can be added')
      log.info('Adding new evidence: {0:s}'.format(str(evidence_)))
      job_count = 0
      jobs_allowlist = evidence_.config.get('jobs_allowlist', [])
      jobs_denylist = evidence_.config.get('jobs_denylist', [])
      if jobs_denylist or jobs_allowlist:
        log.info(
            'Filtering Jobs with allowlist {0!s} and denylist {1!s}'.format(
                jobs_allowlist, jobs_denylist))
        jobs_list = jobs_manager.JobsManager.FilterJobObjects(
            self.jobs, jobs_denylist, jobs_allowlist)
      else:
        jobs_list = self.jobs
  
      # TODO(aarontp): Add some kind of loop detection in here so that jobs can
      # register for Evidence(), or or other evidence types that may be a super
      # class of the output of the job itself.  Short term we could potentially
      # have a run time check for this upon Job instantiation to prevent it.
      for job in jobs_list:
        # Doing a strict type check here for now until we can get the above
        # comment figured out.
        # pylint: disable=unidiomatic-typecheck
        if [True for t in job.evidence_input if type(evidence_) == t]:
          job_instance = job(
              request_id=evidence_.request_id, evidence_config=evidence_.config)
          self.running_jobs.append(job_instance)
          log.info(
              'Adding {0:s} job to process {1:s}'.format(
                  job_instance.name, evidence_.name))
          job_count += 1
          for task in job_instance.create_tasks([evidence_]):
            self.add_task(task, job_instance, evidence_)
  
      if not job_count:
        log.warning(
            'No Jobs/Tasks were created for Evidence [{0:s}]. '
            'Jobs may need to be configured to allow this type of '
            'Evidence as input'.format(str(evidence_)))
  
    def check_done(self):
      """Checks if we have any outstanding tasks.
  
      Returns:
        bool: Indicating whether we are done.
      """
      return not bool(len(self.tasks))
  
    def check_request_done(self, request_id):
      """Checks if we have any outstanding tasks for the request ID.
  
      Args:
        request_id (str): The request ID to check for completion
  
      Returns:
        bool: Indicating whether all Jobs are done.
      """
      job_completion = []
      for job in self.running_jobs:
        if request_id == job.request_id:
          job_completion.append(job.check_done())
  
      return min(job_completion)
  
    def check_request_finalized(self, request_id):
      """Checks if the the request is done and finalized.
  
      A request can be done but not finalized if all of the Tasks created by the
      original Jobs have completed, but the "finalize" Job/Tasks have not been
      run.  These finalize Job/Tasks are created after all of the original
      Jobs/Tasks have completed. Only one Job needs to be marked as finalized for
      the entire request to be considered finalized.
  
      Args:
        request_id (str): The request ID to check for finalization.
  
      Returns:
        bool: Indicating whether all Jobs are done.
      """
      request_finalized = False
      for job in self.running_jobs:
        if request_id == job.request_id and job.is_finalized:
          request_finalized = True
          break
  
      return request_finalized and self.check_request_done(request_id)
  
    def get_evidence(self):
      """Checks for new evidence to process.
  
      Returns:
        list[evidence.Evidence]: The evidence to process.
      """
      raise NotImplementedError
  
    def get_job(self, job_id):
      """Gets the running Job instance from the given Job ID
  
      Args:
        job_id (str): The Job id to get the job for.
  
      Returns:
        TurbiniaJob|None: Job instance if found, else None
      """
      job = None
      for job_instance in self.running_jobs:
        if job_id == job_instance.id:
          job = job_instance
          break
  
      return job
  
    def generate_request_finalize_tasks(self, job):
      """Generates the Tasks to finalize the given request ID.
  
      Args:
        job (TurbiniaJob): The last Job that was run for this request.
      """
      request_id = job.request_id
      final_job = jobs_manager.JobsManager.GetJobInstance('FinalizeRequestJob')
      final_job.request_id = request_id
      final_job.evidence.config = job.evidence.config
      log.debug(
          'Request {0:s} done, but not finalized, creating FinalizeRequestJob '
          '{1:s}'.format(request_id, final_job.id))
  
      # Finalize tasks use EvidenceCollection with all evidence created by the
      # request or job.
      final_evidence = evidence.EvidenceCollection()
      final_evidence.request_id = request_id
      self.running_jobs.append(final_job)
  
      # Gather evidence created by every Job in the request.
      for running_job in self.running_jobs:
        if running_job.request_id == request_id:
          final_evidence.collection.extend(running_job.evidence.collection)
  
      for finalize_task in final_job.create_tasks([final_evidence]):
        self.add_task(finalize_task, final_job, final_evidence)
  
    def add_task(self, task, job, evidence_):
      """Adds a task and evidence to process to the task manager.
  
      Args:
        task: An instantiated Turbinia Task
        evidence_: An Evidence object to be processed.
      """
      if evidence_.request_id:
        task.request_id = evidence_.request_id
      elif job and job.request_id:
        task.request_id = job.request_id
      else:
        log.error(
            'Request ID not found in Evidence {0!s} or Task {1!s}. Not adding '
            'new Task because of undefined state'.format(evidence_, task))
        return
  
      evidence_.config = job.evidence.config
      task.base_output_dir = config.OUTPUT_DIR
      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_)
  
    def remove_jobs(self, request_id):
      """Removes the all Jobs for the given request ID.
  
      Args:
        request_id (str): The ID of the request we want to remove jobs for.
      """
      remove_jobs = [j for j in self.running_jobs if j.request_id == request_id]
      log.debug(
          'Removing {0:d} completed Job(s) for request ID {1:s}.'.format(
              len(remove_jobs), request_id))
      # pylint: disable=expression-not-assigned
      [self.remove_job(j.id) for j in remove_jobs]
  
    def remove_job(self, job_id):
      """Removes a Job from the running jobs list.
  
      Args:
        job_id (str): The ID of the job to remove.
  
      Returns:
        bool: True if Job removed, else False.
      """
      remove_job = None
      for job in self.running_jobs:
        if job_id == job.id:
          remove_job = job
          break
  
      if remove_job:
        self.running_jobs.remove(remove_job)
      return bool(remove_job)
  
    def enqueue_task(self, task, evidence_):
      """Enqueues a task and evidence in the implementation specific task queue.
  
      Args:
        task: An instantiated Turbinia Task
        evidence_: An Evidence object to be processed.
      """
      raise NotImplementedError
  
    def process_result(self, task_result):
      """Runs final task results recording.
  
      self.process_tasks handles things that have failed at the task queue layer
      (i.e. PSQ), and this method handles tasks that have potentially failed
      below that layer (i.e. somewhere in our Task code).
  
      This also adds the Evidence to the running jobs and running requests so we
      can process those later in 'finalize' Tasks.
  
      Args:
        task_result: The TurbiniaTaskResult object
  
      Returns:
        TurbiniaJob|None: The Job for the processed task, else None
      """
      if not task_result.successful:
        log.error(
            'Task {0:s} from {1:s} was not successful'.format(
                task_result.task_name, task_result.worker_name))
      else:
        log.info(
            'Task {0:s} from {1:s} executed with status [{2:s}]'.format(
                task_result.task_name, task_result.worker_name,
                task_result.status))
  
      if not isinstance(task_result.evidence, list):
        log.warning(
            'Task {0:s} from {1:s} did not return evidence list'.format(
                task_result.task_name, task_result.worker_name))
        task_result.evidence = []
  
      job = self.get_job(task_result.job_id)
      if not job:
        log.warning(
            'Received task results for unknown Job from Task ID {0:s}'.format(
                task_result.task_id))
  
      # Reprocess new evidence and save instance for later consumption by finalize
      # tasks.
      for evidence_ in task_result.evidence:
        if isinstance(evidence_, evidence.Evidence):
          log.info(
              'Task {0:s} from {1:s} returned Evidence {2:s}'.format(
                  task_result.task_name, task_result.worker_name, evidence_.name))
          self.add_evidence(evidence_)
          if job:
            job.evidence.add_evidence(evidence_)
        else:
          log.error(
              'Task {0:s} from {1:s} returned non-Evidence output type '
              '{2:s}'.format(
                  task_result.task_name, task_result.worker_name,
                  type(task_result.evidence)))
  
      return job
  
    def process_job(self, job, task):
      """Processes the Job after Task completes.
  
      This removes the Task from the running Job and generates the "finalize"
      Tasks after all the Tasks for the Job and Request have completed.  It also
      removes all Jobs from the running Job list once everything is complete.
  
      Args:
        job (TurbiniaJob): The Job to process
        task (TurbiniaTask): The Task that just completed.
      """
      log.debug(
          'Processing Job {0:s} for completed Task {1:s}'.format(
              job.name, task.id))
      self.state_manager.update_task(task)
      job.remove_task(task.id)
      if job.check_done() and not (job.is_finalize_job or task.is_finalize_task):
        log.debug(
            'Job {0:s} completed, creating Job finalize tasks'.format(job.name))
        final_task = job.create_final_task()
        if final_task:
          final_task.is_finalize_task = True
          self.add_task(final_task, job, job.evidence)
      elif job.check_done() and job.is_finalize_job:
        job.is_finalized = True
  
      request_id = job.request_id
      request_done = self.check_request_done(request_id)
      request_finalized = self.check_request_finalized(request_id)
      # If the request is done but not finalized, we generate the finalize tasks.
      if request_done and not request_finalized:
        self.generate_request_finalize_tasks(job)
  
      # If the Job has been finalized then we can remove all the Jobs for this
      # request since everything is complete.
      elif request_done and request_finalized:
        self.remove_jobs(request_id)
  
    def process_tasks(self):
      """Process any tasks that need to be processed.
  
      Returns:
        list[TurbiniaTask]: Tasks to process that have completed.
      """
      raise NotImplementedError
  
    def run(self, under_test=False):
      """Main run loop for TaskManager."""
      log.info('Starting Task Manager run loop')
      while True:
        # pylint: disable=expression-not-assigned
        [self.add_evidence(x) for x in self.get_evidence()]
  
        for task in self.process_tasks():
          if task.result:
            job = self.process_result(task.result)
            if job:
              self.process_job(job, task)
  
        [self.state_manager.update_task(t) for t in self.tasks]
        if config.SINGLE_RUN and self.check_done():
          log.info('No more tasks to process.  Exiting now.')
          return
  
        if under_test:
          break
  
        time.sleep(config.SLEEP_TIME)
  
  
  class CeleryTaskManager(BaseTaskManager):
    """Celery implementation of BaseTaskManager.
  
    Attributes:
      celery (TurbiniaCelery): Celery task queue, handles worker tasks.
      kombu (TurbiniaKombu): Kombu queue, handles receiving evidence.
      celery_runner: task_runner method, but wrapped for Celery usage.
    """
  
    def __init__(self):
      self.celery = None
      self.kombu = None
      self.celery_runner = None
      config.LoadConfig()
      super(CeleryTaskManager, self).__init__()
  
    def _backend_setup(self, *args, **kwargs):
      self.celery = turbinia_celery.TurbiniaCelery()
      self.celery.setup()
      self.kombu = turbinia_celery.TurbiniaKombu(config.KOMBU_CHANNEL)
      self.kombu.setup()
      self.celery_runner = self.celery.app.task(task_runner, name="task_runner")
  
    def process_tasks(self):
      """Determine the current state of our tasks.
  
      Returns:
        list[TurbiniaTask]: all completed tasks
      """
      completed_tasks = []
      for task in self.tasks:
        celery_task = task.stub
        if not celery_task:
          log.debug('Task {0:s} not yet created'.format(task.stub.task_id))
        elif celery_task.status == celery_states.STARTED:
          log.debug('Task {0:s} not finished'.format(celery_task.id))
        elif celery_task.status == celery_states.FAILURE:
          log.warning('Task {0:s} failed.'.format(celery_task.id))
          completed_tasks.append(task)
        elif celery_task.status == celery_states.SUCCESS:
          task.result = workers.TurbiniaTaskResult.deserialize(celery_task.result)
          completed_tasks.append(task)
        else:
          log.debug('Task {0:s} status unknown'.format(celery_task.id))
  
      outstanding_task_count = len(self.tasks) - len(completed_tasks)
      log.info('{0:d} Tasks still outstanding.'.format(outstanding_task_count))
      return completed_tasks
  
    def get_evidence(self):
      """Receives new evidence.
  
      Returns:
        list[Evidence]: evidence to process.
      """
      requests = self.kombu.check_messages()
      evidence_list = []
      for request in requests:
        for evidence_ in request.evidence:
          if not evidence_.request_id:
            evidence_.request_id = request.request_id
          evidence_.config = request.recipe
          evidence_.config['requester'] = request.requester
          log.info(
              'Received evidence [{0:s}] from Kombu message.'.format(
                  str(evidence_)))
          evidence_list.append(evidence_)
      return evidence_list
  
    def enqueue_task(self, task, evidence_):
      log.info(
          'Adding Celery task {0:s} with evidence {1:s} to queue'.format(
              task.name, evidence_.name))
      task.stub = self.celery_runner.delay(
          task.serialize(), evidence_.serialize())
  
  
  class PSQTaskManager(BaseTaskManager):
    """PSQ implementation of BaseTaskManager.
  
    Attributes:
      psq: PSQ Queue object.
      server_pubsub: A PubSubClient object for receiving new evidence messages.
    """
  
    def __init__(self):
      self.psq = None
      self.server_pubsub = None
      config.LoadConfig()
      super(PSQTaskManager, self).__init__()
  
    # pylint: disable=keyword-arg-before-vararg
    def _backend_setup(self, server=True, *args, **kwargs):
      """
      Args:
        server (bool): Whether this is the client or a server
  
      Raises:
        TurbiniaException: When there are errors creating PSQ Queue
      """
  
      log.debug(
          'Setting up PSQ Task Manager requirements on project {0:s}'.format(
              config.TURBINIA_PROJECT))
      self.server_pubsub = turbinia_pubsub.TurbiniaPubSub(config.PUBSUB_TOPIC)
      if server:
        self.server_pubsub.setup_subscriber()
      else:
        self.server_pubsub.setup_publisher()
      psq_publisher = pubsub.PublisherClient()
      psq_subscriber = pubsub.SubscriberClient()
      datastore_client = datastore.Client(project=config.TURBINIA_PROJECT)
      try:
        self.psq = psq.Queue(
            psq_publisher, psq_subscriber, config.TURBINIA_PROJECT,
            name=config.PSQ_TOPIC, storage=psq.DatastoreStorage(datastore_client))
      except exceptions.GoogleCloudError as e:
        msg = 'Error creating PSQ Queue: {0:s}'.format(str(e))
        log.error(msg)
        raise turbinia.TurbiniaException(msg)
  
    def process_tasks(self):
      completed_tasks = []
      for task in self.tasks:
        psq_task = task.stub.get_task()
        # This handles tasks that have failed at the PSQ layer.
        if not psq_task:
          log.debug('Task {0:s} not yet created'.format(task.stub.task_id))
        elif psq_task.status not in (psq.task.FINISHED, psq.task.FAILED):
          log.debug('Task {0:s} not finished'.format(psq_task.id))
        elif psq_task.status == psq.task.FAILED:
          log.warning('Task {0:s} failed.'.format(psq_task.id))
          completed_tasks.append(task)
        else:
          task.result = workers.TurbiniaTaskResult.deserialize(
              task.stub.result(timeout=PSQ_TASK_TIMEOUT_SECONDS))
          completed_tasks.append(task)
  
      outstanding_task_count = len(self.tasks) - len(completed_tasks)
      log.info('{0:d} Tasks still outstanding.'.format(outstanding_task_count))
      return completed_tasks
  
    def get_evidence(self):
      requests = self.server_pubsub.check_messages()
      evidence_list = []
      for request in requests:
        for evidence_ in request.evidence:
          if not evidence_.request_id:
            evidence_.request_id = request.request_id
          evidence_.config = request.recipe
          evidence_.config['requester'] = request.requester
          log.info(
              'Received evidence [{0:s}] from PubSub message.'.format(
                  str(evidence_)))
          evidence_list.append(evidence_)
      return evidence_list
  
    def enqueue_task(self, task, evidence_):
      log.info(
          'Adding PSQ task {0:s} with evidence {1:s} to queue'.format(
              task.name, evidence_.name))
      task.stub = self.psq.enqueue(
          task_runner, task.serialize(), evidence_.serialize())
      time.sleep(PSQ_QUEUE_WAIT_SECONDS)
  