<commit_before># 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

import logging
import time
import traceback

import psq
from google.cloud import datastore
from google.cloud import pubsub
from google.gax.errors import GaxError

import turbinia
from turbinia import evidence
from turbinia import config
from turbinia import jobs
from turbinia import pubsub as turbinia_pubsub
from turbinia import state_manager

log = logging.getLogger('turbinia')


def get_task_manager():
  """Return task manager object based on config.

  Returns
    Initialized TaskManager object.
  """
  config.LoadConfig()
  if config.TASK_MANAGER == 'PSQ':
    return PSQTaskManager()
  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).
  """
  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: A list of instantiated job objects
    evidence: A list of evidence objects to process
    state_manager: State manager object to handle syncing with storage
    tasks: A list of outstanding TurbiniaTask objects
  """

  def __init__(self):
    self.jobs = []
    self.evidence = []
    self.tasks = []
    self.state_manager = state_manager.get_state_manager()

  def _backend_setup(self):
    """Sets up backend dependencies.

    Raises:
      TurbiniaException: When encountering fatal errors setting up dependencies.
    """
    raise NotImplementedError

  def setup(self):
    """Does setup of Task manager and its dependencies."""
    self._backend_setup()
    # TODO(aarontp): Consider instantiating a job per evidence object
    self.jobs = jobs.get_jobs()

  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_)))
    self.evidence.append(evidence_)
    job_count = 0
    # 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 self.jobs:
      if [True for t in job.evidence_input if isinstance(evidence_, t)]:
        log.info(
            'Adding {0:s} job to process {1:s}'.format(
                job.name, evidence_.name))
        job_count += 1
        for task in job.create_tasks([evidence_]):
          task.base_output_dir = config.OUTPUT_DIR
          self.add_task(task, 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 to see if we have any outstanding tasks.<commit_msg>Strict type check when matching evidence to jobs (#123)<commit_after># 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

import logging
import time
import traceback

import psq
from google.cloud import datastore
from google.cloud import pubsub
from google.gax.errors import GaxError

import turbinia
from turbinia import evidence
from turbinia import config
from turbinia import jobs
from turbinia import pubsub as turbinia_pubsub
from turbinia import state_manager

log = logging.getLogger('turbinia')


def get_task_manager():
  """Return task manager object based on config.

  Returns
    Initialized TaskManager object.
  """
  config.LoadConfig()
  if config.TASK_MANAGER == 'PSQ':
    return PSQTaskManager()
  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).
  """
  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: A list of instantiated job objects
    evidence: A list of evidence objects to process
    state_manager: State manager object to handle syncing with storage
    tasks: A list of outstanding TurbiniaTask objects
  """

  def __init__(self):
    self.jobs = []
    self.evidence = []
    self.tasks = []
    self.state_manager = state_manager.get_state_manager()

  def _backend_setup(self):
    """Sets up backend dependencies.

    Raises:
      TurbiniaException: When encountering fatal errors setting up dependencies.
    """
    raise NotImplementedError

  def setup(self):
    """Does setup of Task manager and its dependencies."""
    self._backend_setup()
    # TODO(aarontp): Consider instantiating a job per evidence object
    self.jobs = jobs.get_jobs()

  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_)))
    self.evidence.append(evidence_)
    job_count = 0
    # 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 self.jobs:
      # 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]:
        log.info(
            'Adding {0:s} job to process {1:s}'.format(
                job.name, evidence_.name))
        job_count += 1
        for task in job.create_tasks([evidence_]):
          task.base_output_dir = config.OUTPUT_DIR
          self.add_task(task, 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 to see if we have any outstanding tasks.
{'fraction_non_alphanumeric': 0.076449812008659, 'fraction_numerical': 0.0037598268200979834, 'mean_word_length': 3.3241379310344827, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3678284', 'n_tokens_mistral': 2541, 'n_tokens_neox': 2387, 'n_words': 1006}