"""The task-based multi-process processing analysis engine."""

import collections
import os
import time

from plaso.containers import counts
from plaso.containers import events
from plaso.containers import reports
from plaso.containers import tasks
from plaso.engine import processing_status
from plaso.lib import definitions
from plaso.lib import errors
from plaso.multi_process import analysis_process
from plaso.multi_process import logger
from plaso.multi_process import merge_helpers
from plaso.multi_process import plaso_queue
from plaso.multi_process import task_engine
from plaso.multi_process import zeromq_queue


class AnalysisMultiProcessEngine(task_engine.TaskMultiProcessEngine):
  """Task-based multi-process analysis engine.

  This class contains functionality to:
  * monitor and manage analysis tasks;
  * merge results returned by analysis worker processes.
  """

  # pylint: disable=abstract-method

  _CONTAINER_TYPE_ANALYSIS_REPORT = reports.AnalysisReport.CONTAINER_TYPE
  _CONTAINER_TYPE_EVENT_TAG = events.EventTag.CONTAINER_TYPE

  _PROCESS_JOIN_TIMEOUT = 5.0

  _QUEUE_TIMEOUT = 10 * 60

  def __init__(self, worker_memory_limit=None, worker_timeout=None):
    """Initializes a task-based multi-process analysis engine.

    Args:
      worker_memory_limit (Optional[int]): maximum amount of memory a worker is
          allowed to consume, where None represents the default memory limit
          and 0 represents no limit.
      worker_timeout (Optional[float]): number of minutes before a worker
          process that is not providing status updates is considered inactive,
          where None or 0.0 represents the default timeout.
    """
    if worker_memory_limit is None:
      worker_memory_limit = definitions.DEFAULT_WORKER_MEMORY_LIMIT

    super().__init__()
    self._analysis_plugins = {}
    self._completed_analysis_processes = set()
    self._data_location = None
    self._event_filter_expression = None
    self._event_labels_counter = None
    self._event_queues = {}
    self._events_status = processing_status.EventsStatus()
    self._memory_profiler = None
    self._merge_task = None
    self._number_of_consumed_analysis_reports = 0
    self._number_of_consumed_event_data = 0
    self._number_of_consumed_event_tags = 0
    self._number_of_consumed_events = 0
    self._number_of_consumed_sources = 0
    self._number_of_produced_analysis_reports = 0
    self._number_of_produced_event_data = 0
    self._number_of_produced_event_tags = 0
    self._number_of_produced_events = 0
    self._number_of_produced_sources = 0
    self._processing_profiler = None
    self._serializers_profiler = None
    self._session = None
    self._status = definitions.STATUS_INDICATOR_IDLE
    self._status_update_callback = None
    self._user_accounts = None
    self._worker_memory_limit = worker_memory_limit
    self._worker_timeout = worker_timeout or definitions.DEFAULT_WORKER_TIMEOUT

  def _AnalyzeEvents(self, storage_writer, analysis_plugins, event_filter=None):
    """Analyzes events in a Plaso storage.

    Args:
      storage_writer (StorageWriter): storage writer.
      analysis_plugins (dict[str, AnalysisPlugin]): analysis plugins that
          should be run and their names.
      event_filter (Optional[EventObjectFilter]): event filter.

    Returns:
      collections.Counter: counter containing information about the events
          processed and filtered.

    Raises:
      RuntimeError: if a non-recoverable situation is encountered.
    """
    self._status = definitions.STATUS_INDICATOR_RUNNING
    self._number_of_consumed_analysis_reports = 0
    self._number_of_consumed_event_data = 0
    self._number_of_consumed_event_tags = 0
    self._number_of_consumed_events = 0
    self._number_of_consumed_sources = 0
    self._number_of_produced_analysis_reports = 0
    self._number_of_produced_event_data = 0
    self._number_of_produced_event_tags = 0
    self._number_of_produced_events = 0
    self._number_of_produced_sources = 0

    number_of_filtered_events = 0

    logger.debug('Processing events.')

    filter_limit = getattr(event_filter, 'limit', None)

    for event in storage_writer.GetSortedEvents():
      event_data_identifier = event.GetEventDataIdentifier()
      event_data = storage_writer.GetAttributeContainerByIdentifier(
          events.EventData.CONTAINER_TYPE, event_data_identifier)

      event_data_stream_identifier = event_data.GetEventDataStreamIdentifier()
      if event_data_stream_identifier:
        event_data_stream = storage_writer.GetAttributeContainerByIdentifier(
            events.EventDataStream.CONTAINER_TYPE, event_data_stream_identifier)
      else:
        event_data_stream = None

      event_identifier = event.GetIdentifier()
      event_tag = storage_writer.GetEventTagByEventIdentifer(event_identifier)

      if event_filter:
        filter_match = event_filter.Match(
            event, event_data, event_data_stream, event_tag)
      else:
        filter_match = None

      # pylint: disable=singleton-comparison
      if filter_match == False:
        number_of_filtered_events += 1
        continue

      for event_queue in self._event_queues.values():
        # TODO: Check for premature exit of analysis plugins.
        event_tripple = events.EventTripple()
        event_tripple.event = event
        event_tripple.event_data = event_data
        event_tripple.event_data_stream = event_data_stream

        event_queue.PushItem(event_tripple)

      self._number_of_consumed_events += 1

      if (event_filter and filter_limit and
          filter_limit == self._number_of_consumed_events):
        break

    logger.debug('Finished pushing events to analysis plugins.')
    # Signal that we have finished adding events.
    for event_queue in self._event_queues.values():
      event_queue.PushItem(plaso_queue.QueueAbort(), block=False)

    logger.debug('Processing analysis plugin results.')

    # TODO: use a task based approach.
    plugin_names = list(analysis_plugins.keys())
    while plugin_names:
      for plugin_name in list(plugin_names):
        if self._abort:
          break

        # TODO: temporary solution.
        task = tasks.Task()
        task.storage_format = definitions.STORAGE_FORMAT_SQLITE
        task.identifier = plugin_name

        merge_ready = self._CheckTaskReadyForMerge(
            definitions.STORAGE_FORMAT_SQLITE, task)
        if merge_ready:
          self._PrepareMergeTaskStorage(definitions.STORAGE_FORMAT_SQLITE, task)
          self._status = definitions.STATUS_INDICATOR_MERGING

          event_queue = self._event_queues[plugin_name]
          del self._event_queues[plugin_name]

          event_queue.Close()

          task_storage_reader = self._GetMergeTaskStorage(
              definitions.STORAGE_FORMAT_SQLITE, task)

          try:
            merge_helper = merge_helpers.AnalysisTaskMergeHelper(
                task_storage_reader, task.identifier)

            logger.debug(
                f'Starting merge of task: {merge_helper.task_identifier:s}')

            number_of_containers = self._MergeAttributeContainers(
                storage_writer, merge_helper)

            logger.debug((
                f'Merged {number_of_containers:d} containers of task: '
                f'{merge_helper.task_identifier:s}'))

          finally:
            task_storage_reader.Close()

          self._RemoveMergeTaskStorage(
              definitions.STORAGE_FORMAT_SQLITE, task)

          self._status = definitions.STATUS_INDICATOR_RUNNING

          # TODO: temporary solution.
          plugin_names.remove(plugin_name)

    events_counter = collections.Counter()
    events_counter['Events filtered'] = number_of_filtered_events
    events_counter['Events processed'] = self._number_of_consumed_events

    return events_counter

  def _CheckStatusAnalysisProcess(self, pid):
    """Checks the status of an analysis process.

    Args:
      pid (int): process ID (PID) of a registered analysis process.

    Raises:
      KeyError: if the process is not registered with the engine.
    """
    # TODO: Refactor this method, simplify and separate concerns (monitoring
    # vs management).
    self._RaiseIfNotRegistered(pid)

    if pid in self._completed_analysis_processes:
      status_indicator = definitions.STATUS_INDICATOR_COMPLETED
      process_status = {
          'processing_status': status_indicator}
      used_memory = 0

    else:
      process = self._processes_per_pid[pid]

      process_status = self._QueryProcessStatus(process)
      if process_status is None:
        process_is_alive = False
      else:
        process_is_alive = True

      process_information = self._process_information_per_pid[pid]
      used_memory = process_information.GetUsedMemory() or 0

      if self._worker_memory_limit and used_memory > self._worker_memory_limit:
        logger.warning((
            f'Process: {process.name:s} (PID: {pid:d}) killed because it '
            f'exceeded the memory limit: {self._worker_memory_limit:d}'))
        self._KillProcess(pid)

      if isinstance(process_status, dict):
        self._rpc_errors_per_pid[pid] = 0
        status_indicator = process_status.get('processing_status')

        if status_indicator == definitions.STATUS_INDICATOR_COMPLETED:
          self._completed_analysis_processes.add(pid)

      else:
        rpc_errors = self._rpc_errors_per_pid.get(pid, 0) + 1
        self._rpc_errors_per_pid[pid] = rpc_errors

        if rpc_errors > self._MAXIMUM_RPC_ERRORS:
          process_is_alive = False

        if process_is_alive:
          rpc_port = process.rpc_port.value
          logger.warning((
              f'Unable to retrieve process: {process.name:s} (PID: {pid:d}) '
              f'status via RPC socket: http://localhost:{rpc_port:d}'))

          processing_status_string = 'RPC error'
          status_indicator = definitions.STATUS_INDICATOR_RUNNING
        else:
          processing_status_string = 'killed'
          status_indicator = definitions.STATUS_INDICATOR_KILLED

        process_status = {
            'processing_status': processing_status_string}

    self._UpdateProcessingStatus(pid, process_status, used_memory)

    if status_indicator in definitions.ERROR_STATUS_INDICATORS:
      logger.error((
          f'Process {process.name:s} (PID: {pid:d}) is not functioning '
          f'correctly. Status code: {status_indicator!s}.'))

      self._TerminateProcessByPid(pid)

  def _MergeAttributeContainers(self, storage_writer, merge_helper):
    """Merges attribute containers from a task store into the storage writer.

    Args:
      storage_writer (StorageWriter): storage writer.
      merge_helper (AnalysisTaskMergeHelper): helper to merge attribute
          containers.

    Returns:
      int: number of containers merged.
    """
    number_of_containers = 0

    container = merge_helper.GetAttributeContainer()

    while container:
      number_of_containers += 1

      if container.CONTAINER_TYPE == self._CONTAINER_TYPE_EVENT_TAG:
        storage_writer.AddOrUpdateEventTag(container)
      else:
        storage_writer.AddAttributeContainer(container)

      if container.CONTAINER_TYPE == self._CONTAINER_TYPE_ANALYSIS_REPORT:
        self._number_of_produced_analysis_reports += 1

      elif container.CONTAINER_TYPE == self._CONTAINER_TYPE_EVENT_TAG:
        self._number_of_produced_event_tags += 1

        for label in container.labels:
          self._event_labels_counter[label] += 1
          self._event_labels_counter['total'] += 1

      container = merge_helper.GetAttributeContainer()

    return number_of_containers

  def _StartAnalysisProcesses(self, analysis_plugins):
    """Starts the analysis processes.

    Args:
      analysis_plugins (dict[str, AnalysisPlugin]): analysis plugins that
          should be run and their names.
    """
    logger.info('Starting analysis plugins.')

    for analysis_plugin in analysis_plugins.values():
      self._analysis_plugins[analysis_plugin.NAME] = analysis_plugin

      process = self._StartWorkerProcess(analysis_plugin.NAME)
      if not process:
        logger.error(
            f'Unable to create analysis process: {analysis_plugin.NAME:s}')

    logger.info('Analysis plugins running')

  def _StartWorkerProcess(self, process_name):
    """Creates, starts, monitors and registers a worker process.

    Args:
      process_name (str): process name.

    Returns:
      MultiProcessWorkerProcess: extraction worker process or None on error.
    """
    analysis_plugin = self._analysis_plugins.get(process_name)
    if not analysis_plugin:
      logger.error(f'Missing analysis plugin: {process_name:s}')
      return None

    queue_name = f'{process_name:s} output event queue'
    output_event_queue = zeromq_queue.ZeroMQPushBindQueue(
        name=queue_name, timeout_seconds=self._QUEUE_TIMEOUT)
    # Open the queue so it can bind to a random port, and we can get the
    # port number to use in the input queue.
    output_event_queue.Open()

    self._event_queues[process_name] = output_event_queue

    queue_name = f'{process_name:s} input event queue'
    input_event_queue = zeromq_queue.ZeroMQPullConnectQueue(
        name=queue_name, delay_open=True, port=output_event_queue.port,
        timeout_seconds=self._QUEUE_TIMEOUT)

    process = analysis_process.AnalysisProcess(
        input_event_queue, analysis_plugin, self._processing_configuration,
        self._user_accounts, data_location=self._data_location,
        event_filter_expression=self._event_filter_expression,
        name=process_name)

    process.start()

    logger.info(
        f'Started analysis plugin: {process_name:s} (PID: {process.pid:d}).')

    try:
      self._StartMonitoringProcess(process)
    except (KeyError, OSError) as exception:
      logger.error((
          f'Unable to monitor analysis plugin: {process_name:s} (PID: '
          f'{process.pid:d}) with error: {exception!s}'))

      process.terminate()
      return None

    self._RegisterProcess(process)
    return process

  def _StopAnalysisProcesses(self, abort=False):
    """Stops the analysis processes.

    Args:
      abort (bool): True to indicated the stop is issued on abort.
    """
    logger.debug('Stopping analysis processes.')
    self._StopMonitoringProcesses()

    if abort:
      # Signal all the processes to abort.
      self._AbortTerminate()

    # Wake the processes to make sure that they are not blocking
    # waiting for the queue new items.
    for event_queue in self._event_queues.values():
      event_queue.PushItem(plaso_queue.QueueAbort(), block=False)

    # Try waiting for the processes to exit normally.
    self._AbortJoin(timeout=self._PROCESS_JOIN_TIMEOUT)
    for event_queue in self._event_queues.values():
      event_queue.Close(abort=abort)

    if abort:
      # Kill any remaining processes.
      self._AbortKill()
    else:
      # Check if the processes are still alive and terminate them if necessary.
      self._AbortTerminate()
      self._AbortJoin(timeout=self._PROCESS_JOIN_TIMEOUT)

      for event_queue in self._event_queues.values():
        event_queue.Close(abort=True)

  def _UpdateForemanProcessStatus(self):
    """Update the foreman process status."""
    used_memory = self._process_information.GetUsedMemory() or 0

    display_name = getattr(self._merge_task, 'identifier', '')

    self._processing_status.UpdateForemanStatus(
        self._name, self._status, self._pid, used_memory, display_name,
        self._number_of_consumed_sources, self._number_of_produced_sources,
        self._number_of_consumed_event_data,
        self._number_of_produced_event_data,
        self._number_of_consumed_events, self._number_of_produced_events,
        self._number_of_consumed_event_tags,
        self._number_of_produced_event_tags,
        self._number_of_consumed_analysis_reports,
        self._number_of_produced_analysis_reports)

    self._processing_status.UpdateEventsStatus(self._events_status)

  def _UpdateProcessingStatus(self, pid, process_status, used_memory):
    """Updates the processing status.

    Args:
      pid (int): process identifier (PID) of the worker process.
      process_status (dict[str, object]): status values received from
          the worker process.
      used_memory (int): size of used memory in bytes.

    Raises:
      KeyError: if the process is not registered with the engine.
    """
    self._RaiseIfNotRegistered(pid)

    if not process_status:
      return

    process = self._processes_per_pid[pid]

    status_indicator = process_status.get('processing_status')

    self._RaiseIfNotMonitored(pid)

    display_name = process_status.get('display_name', '')

    number_of_consumed_event_data = process_status.get(
        'number_of_consumed_event_data', None)
    number_of_produced_event_data = process_status.get(
        'number_of_produced_event_data', None)

    number_of_consumed_event_tags = process_status.get(
        'number_of_consumed_event_tags', None)
    number_of_produced_event_tags = process_status.get(
        'number_of_produced_event_tags', None)

    number_of_consumed_events = process_status.get(
        'number_of_consumed_events', None)
    number_of_produced_events = process_status.get(
        'number_of_produced_events', None)

    number_of_consumed_reports = process_status.get(
        'number_of_consumed_reports', None)
    number_of_produced_reports = process_status.get(
        'number_of_produced_reports', None)

    number_of_consumed_sources = process_status.get(
        'number_of_consumed_sources', None)
    number_of_produced_sources = process_status.get(
        'number_of_produced_sources', None)

    if status_indicator != definitions.STATUS_INDICATOR_IDLE:
      last_activity_timestamp = process_status.get(
          'last_activity_timestamp', 0.0)

      if last_activity_timestamp:
        last_activity_timestamp += self._worker_timeout

        current_timestamp = time.time()
        if current_timestamp > last_activity_timestamp:
          logger.error((
              f'Process {process.name:s} (PID: {pid:d}) has not reported '
              f'activity within the timeout period.'))
          status_indicator = definitions.STATUS_INDICATOR_NOT_RESPONDING

    self._processing_status.UpdateWorkerStatus(
        process.name, status_indicator, pid, used_memory, display_name,
        number_of_consumed_sources, number_of_produced_sources,
        number_of_consumed_event_data, number_of_produced_event_data,
        number_of_consumed_events, number_of_produced_events,
        number_of_consumed_event_tags, number_of_produced_event_tags,
        number_of_consumed_reports, number_of_produced_reports)

  def _UpdateStatus(self):
    """Update the status."""
    # Make a local copy of the PIDs in case the dict is changed by
    # the main thread.
    for pid in list(self._process_information_per_pid.keys()):
      self._CheckStatusAnalysisProcess(pid)

    self._UpdateForemanProcessStatus()

    if self._status_update_callback:
      self._status_update_callback(self._processing_status)

  # pylint: disable=too-many-arguments
  def AnalyzeEvents(
      self, session, storage_writer, data_location, analysis_plugins,
      processing_configuration, event_filter=None,
      event_filter_expression=None, status_update_callback=None,
      storage_file_path=None):
    """Analyzes events in a Plaso storage.

    Args:
      session (Session): session in which the events are analyzed.
      storage_writer (StorageWriter): storage writer.
      data_location (str): path to the location that data files should
          be loaded from.
      analysis_plugins (dict[str, AnalysisPlugin]): analysis plugins that
          should be run and their names.
      processing_configuration (ProcessingConfiguration): processing
          configuration.
      event_filter (Optional[EventObjectFilter]): event filter.
      event_filter_expression (Optional[str]): event filter expression.
      status_update_callback (Optional[function]): callback function for status
          updates.
      storage_file_path (Optional[str]): path to the session storage file.

    Returns:
      ProcessingStatus: processing status.

    Raises:
      KeyboardInterrupt: if a keyboard interrupt was raised.
      ValueError: if analysis plugins are missing.
    """
    if not analysis_plugins:
      raise ValueError('Missing analysis plugins')

    abort_kill = False
    keyboard_interrupt = False
    queue_full = False

    self._analysis_plugins = {}
    self._data_location = data_location
    self._event_filter_expression = event_filter_expression
    self._events_status = processing_status.EventsStatus()
    self._processing_configuration = processing_configuration
    self._session = session
    self._status_update_callback = status_update_callback
    self._storage_file_path = storage_file_path

    self._user_accounts = list(
        storage_writer.GetAttributeContainers('user_account'))

    stored_event_labels_counter = {}
    if storage_writer.HasAttributeContainers('event_label_count'):
      stored_event_labels_counter = {
          event_label_count.label: event_label_count
          for event_label_count in storage_writer.GetAttributeContainers(
              'event_label_count')}

    self._event_labels_counter = collections.Counter()

    if storage_writer.HasAttributeContainers('parser_count'):
      parsers_counter = {
          parser_count.name: parser_count.number_of_events
          for parser_count in storage_writer.GetAttributeContainers(
              'parser_count')}

      total_number_of_events = parsers_counter['total']

    else:
      total_number_of_events = 0
      for stored_session in storage_writer.GetSessions():
        total_number_of_events += stored_session.parsers_counter['total']

    self._events_status.total_number_of_events = total_number_of_events

    # Set up the storage writer before the analysis processes.
    self._StartTaskStorage(definitions.STORAGE_FORMAT_SQLITE)

    self._StartAnalysisProcesses(analysis_plugins)

    self._StartProfiling(self._processing_configuration.profiling)

    # Start the status update thread after open of the storage writer
    # so we don't have to clean up the thread if the open fails.
    self._StartStatusUpdateThread()

    try:
      self._AnalyzeEvents(
          storage_writer, analysis_plugins, event_filter=event_filter)

      for key, value in self._event_labels_counter.items():
        event_label_count = stored_event_labels_counter.get(key)
        if event_label_count:
          event_label_count.number_of_events += value
          storage_writer.UpdateAttributeContainer(event_label_count)
        else:
          event_label_count = counts.EventLabelCount(
              label=key, number_of_events=value)
          storage_writer.AddAttributeContainer(event_label_count)

      self._status = definitions.STATUS_INDICATOR_FINALIZING

    except errors.QueueFull:
      queue_full = True
      self._abort = True

    except KeyboardInterrupt:
      keyboard_interrupt = True
      self._abort = True

    finally:
      self._processing_status.aborted = self._abort
      session.aborted = self._abort

      # Stop the status update thread after close of the storage writer
      # so we include the storage sync to disk in the status updates.
      self._StopStatusUpdateThread()

      self._StopProfiling()

    # Update the status view one last time before the analysis processses are
    # stopped.
    self._UpdateStatus()

    if queue_full:
      # TODO: handle abort on queue full more elegant.
      abort_kill = True
    else:
      try:
        self._StopAnalysisProcesses(abort=self._abort)

      except KeyboardInterrupt:
        keyboard_interrupt = True
        abort_kill = True

    if abort_kill:
      self._AbortKill()

      # The abort can leave the main process unresponsive
      # due to incorrectly finalized IPC.
      self._KillProcess(os.getpid())

    try:
      self._StopTaskStorage(
          definitions.STORAGE_FORMAT_SQLITE, session.identifier,
          abort=self._abort)
    except OSError as exception:
      logger.error(f'Unable to stop task storage with error: {exception!s}')

    if self._abort:
      logger.debug('Analysis aborted.')
    else:
      logger.debug('Analysis completed.')

    # Update the status view one last time.
    self._UpdateStatus()

    # Reset values.
    self._analysis_plugins = {}
    self._data_location = None
    self._event_filter_expression = None
    self._processing_configuration = None
    self._session = None
    self._status_update_callback = None
    self._storage_file_path = None
    self._user_accounts = None

    if keyboard_interrupt:
      raise KeyboardInterrupt

    return self._processing_status
