#!/usr/bin/env python
"""Classes for hunt-related testing."""

import hashlib
import sys


from grr_response_core.lib.rdfvalues import client as rdf_client
from grr_response_core.lib.rdfvalues import client_fs as rdf_client_fs
from grr_response_core.lib.rdfvalues import flows as rdf_flows
from grr_response_core.lib.rdfvalues import paths as rdf_paths
from grr_response_core.lib.rdfvalues import structs as rdf_structs
from grr_response_server import data_store
from grr_response_server import flow
from grr_response_server import foreman
from grr_response_server import foreman_rules
from grr_response_server import hunt
from grr_response_server import output_plugin
from grr_response_server.databases import db
from grr_response_server.flows.general import transfer
from grr_response_server.rdfvalues import flow_objects as rdf_flow_objects
from grr_response_server.rdfvalues import flow_runner as rdf_flow_runner
from grr_response_server.rdfvalues import hunt_objects as rdf_hunt_objects
from grr.test_lib import acl_test_lib
from grr.test_lib import action_mocks
from grr.test_lib import flow_test_lib


class SampleHuntMock(action_mocks.ActionMock):
  """Client mock for sample hunts."""

  def __init__(self,
               failrate=-1,
               data=b"Hello World!",
               user_cpu_time=None,
               system_cpu_time=None,
               network_bytes_sent=None):
    super().__init__()
    self.responses = 0
    self.data = data
    self.failrate = failrate
    self.count = 0

    self.user_cpu_time = user_cpu_time
    self.system_cpu_time = system_cpu_time
    self.network_bytes_sent = network_bytes_sent

  def GetFileStat(self, args):
    """`GetFileStat` action mock."""
    response = rdf_client_fs.StatEntry(
        pathspec=args.pathspec,
        st_mode=33184,
        st_ino=1063090,
        st_dev=64512,
        st_nlink=1,
        st_uid=139592,
        st_gid=5000,
        st_size=len(self.data),
        st_atime=1336469177,
        st_mtime=1336129892,
        st_ctime=1336129892)

    self.responses += 1
    self.count += 1

    # Every "failrate" client does not have this file.
    if self.count == self.failrate:
      self.count = 0
      return []

    return [response]

  def GenerateStatusMessage(self, message, response_id, status=None):
    status = rdf_flows.GrrStatus(
        status=status or rdf_flows.GrrStatus.ReturnedStatus.OK)

    if message.name == "GetFileStat":
      # Create status message to report sample resource usage
      if self.user_cpu_time is None:
        status.cpu_time_used.user_cpu_time = self.responses
      else:
        status.cpu_time_used.user_cpu_time = self.user_cpu_time

      if self.system_cpu_time is None:
        status.cpu_time_used.system_cpu_time = self.responses * 2
      else:
        status.cpu_time_used.system_cpu_time = self.system_cpu_time

      if self.network_bytes_sent is None:
        status.network_bytes_sent = self.responses * 3
      else:
        status.network_bytes_sent = self.network_bytes_sent

    return rdf_flows.GrrMessage(
        session_id=message.session_id,
        name=message.name,
        response_id=response_id,
        request_id=message.request_id,
        payload=status,
        type=rdf_flows.GrrMessage.Type.STATUS)

  def TransferBuffer(self, args):
    """TransferBuffer action mock."""
    response = rdf_client.BufferReference(args)

    offset = min(args.offset, len(self.data))
    sha256 = hashlib.sha256()
    sha256.update(self.data[offset:])
    response.data = sha256.digest()
    response.length = len(self.data[offset:])
    data_store.BLOBS.WriteBlobWithUnknownHash(self.data[offset:])

    return [response]


def TestHuntHelperWithMultipleMocks(client_mocks,
                                    iteration_limit=None,
                                    worker=None):
  """Runs a hunt with a given set of clients mocks.

  Args:
    client_mocks: Dictionary of (client_id->client_mock) pairs. Client mock
      objects are used to handle client actions. Methods names of a client mock
      object correspond to client actions names. For an example of a client mock
      object, see SampleHuntMock.
    iteration_limit: If None, hunt will run until it's finished. Otherwise,
      worker_mock.Next() will be called iteration_limit number of times. Every
      iteration processes worker's message queue. If new messages are sent to
      the queue during the iteration processing, they will be processed on next
      iteration.
    worker: flow_test_lib.TestWorker object to use.

  Returns:
    A number of iterations complete.
  """

  client_mocks = [
      flow_test_lib.MockClient(client_id, client_mock)
      for client_id, client_mock in client_mocks.items()
  ]

  if worker is None:
    rel_db_worker = flow_test_lib.TestWorker()
    data_store.REL_DB.RegisterFlowProcessingHandler(rel_db_worker.ProcessFlow)
  else:
    rel_db_worker = worker

  num_iterations = 0

  try:
    # Run the clients and worker until nothing changes any more.
    while iteration_limit is None or num_iterations < iteration_limit:
      data_store.REL_DB.delegate.WaitUntilNoFlowsToProcess(timeout=10)
      worker_processed = rel_db_worker.ResetProcessedFlows()

      client_processed = 0

      for client_mock in client_mocks:
        client_processed += int(client_mock.Next())

      num_iterations += 1

      if client_processed == 0 and not worker_processed:
        break

  finally:
    if worker is None:
      data_store.REL_DB.UnregisterFlowProcessingHandler(timeout=60)
      rel_db_worker.Shutdown()

  return num_iterations


def TestHuntHelper(client_mock,
                   client_ids,
                   iteration_limit=None,
                   worker=None):
  """Runs a hunt with a given client mock on given clients.

  Args:
    client_mock: Client mock objects are used to handle client actions. Methods
      names of a client mock object correspond to client actions names. For an
      example of a client mock object, see SampleHuntMock.
    client_ids: List of clients ids. Hunt will run on these clients. client_mock
      will be used for every client id.
    iteration_limit: If None, hunt will run until it's finished. Otherwise,
      worker_mock.Next() will be called iteration_limit number of tiems. Every
      iteration processes worker's message queue. If new messages are sent to
      the queue during the iteration processing, they will be processed on next
      iteration.
    worker: flow_test_lib.TestWorker object to use.

  Returns:
    A number of iterations complete.
  """
  return TestHuntHelperWithMultipleMocks(
      dict([(client_id, client_mock) for client_id in client_ids]),
      iteration_limit=iteration_limit,
      worker=worker)


class StandardHuntTestMixin(acl_test_lib.AclTestMixin):
  """Mixin with helper methods for hunt tests."""

  def _CreateForemanClientRuleSet(self):
    return foreman_rules.ForemanClientRuleSet(rules=[
        foreman_rules.ForemanClientRule(
            rule_type=foreman_rules.ForemanClientRule.Type.REGEX,
            regex=foreman_rules.ForemanRegexClientRule(
                field="CLIENT_NAME", attribute_regex="GRR"))
    ])

  def CreateHunt(self,
                 flow_runner_args=None,
                 flow_args=None,
                 client_rule_set=None,
                 original_object=None,
                 client_rate=0,
                 duration=None,
                 creator=None,
                 **kwargs):
    # Only initialize default flow_args value if default flow_runner_args value
    # is to be used.
    if not flow_runner_args:
      flow_args = (
          flow_args or transfer.GetFileArgs(
              pathspec=rdf_paths.PathSpec(
                  path="/tmp/evil.txt",
                  pathtype=rdf_paths.PathSpec.PathType.OS)))

    flow_runner_args = (
        flow_runner_args or
        rdf_flow_runner.FlowRunnerArgs(flow_name=transfer.GetFile.__name__))

    client_rule_set = (client_rule_set or self._CreateForemanClientRuleSet())

    hunt_args = rdf_hunt_objects.HuntArguments.Standard(
        flow_name=flow_runner_args.flow_name,
        flow_args=rdf_structs.AnyValue.Pack(flow_args))

    hunt_obj = rdf_hunt_objects.Hunt(
        creator=creator,
        client_rule_set=client_rule_set,
        original_object=original_object,
        client_rate=client_rate,
        duration=duration,
        args=hunt_args,
        **kwargs)
    hunt.CreateHunt(hunt_obj)

    return hunt_obj.hunt_id

  def StartHunt(self, paused=False, **kwargs):
    hunt_id = self.CreateHunt(**kwargs)
    if not paused:
      hunt.StartHunt(hunt_id)
    return hunt_id

  def FindForemanRules(self, hunt_obj):
    rules = data_store.REL_DB.ReadAllForemanRules()
    return [
        rule for rule in rules
        if hunt_obj is None or rule.hunt_id == hunt_obj.urn.Basename()
    ]

  def AssignTasksToClients(self, client_ids=None, worker=None):
    # Pretend to be the foreman now and dish out hunting jobs to all the
    # clients..
    client_ids = client_ids or self.client_ids
    foreman_obj = foreman.Foreman()

    def Assign():
      for client_id in client_ids:
        foreman_obj.AssignTasksToClient(client_id)

    if worker is None:
      with flow_test_lib.TestWorker():
        Assign()
    else:
      Assign()

  def RunHunt(self,
              client_ids=None,
              iteration_limit=None,
              client_mock=None,
              **mock_kwargs):
    with flow_test_lib.TestWorker() as test_worker:
      self.AssignTasksToClients(client_ids=client_ids, worker=test_worker)

      if client_mock is None:
        client_mock = SampleHuntMock(**mock_kwargs)
      return TestHuntHelper(
          client_mock,
          client_ids or self.client_ids,
          iteration_limit=iteration_limit,
          worker=test_worker)

  def RunHuntWithClientCrashes(self, client_ids):
    with flow_test_lib.TestWorker() as test_worker:
      client_mocks = dict([(client_id, flow_test_lib.CrashClientMock(client_id))
                           for client_id in client_ids])
      self.AssignTasksToClients(client_ids=client_ids, worker=test_worker)
      return TestHuntHelperWithMultipleMocks(client_mocks)

  def _EnsureClientHasHunt(self, client_id, hunt_id):
    try:
      data_store.REL_DB.ReadFlowObject(client_id, hunt_id)
    except db.UnknownFlowError:
      flow_test_lib.StartFlow(
          transfer.GetFile,
          client_id=client_id,
          parent=flow.FlowParent.FromHuntID(hunt_id))

    return hunt_id

  def AddResultsToHunt(self, hunt_id, client_id, values):
    flow_id = self._EnsureClientHasHunt(client_id, hunt_id)

    for value in values:
      data_store.REL_DB.WriteFlowResults([
          rdf_flow_objects.FlowResult(
              client_id=client_id,
              flow_id=flow_id,
              hunt_id=hunt_id,
              payload=value)
      ])

  def AddLogToHunt(self, hunt_id, client_id, message):
    flow_id = self._EnsureClientHasHunt(client_id, hunt_id)

    data_store.REL_DB.WriteFlowLogEntry(
        rdf_flow_objects.FlowLogEntry(
            client_id=client_id,
            flow_id=flow_id,
            hunt_id=hunt_id,
            message=message))

  def AddErrorToHunt(self, hunt_id, client_id, message, backtrace):
    flow_id = self._EnsureClientHasHunt(client_id, hunt_id)

    flow_obj = data_store.REL_DB.ReadFlowObject(client_id, flow_id)
    flow_obj.flow_state = flow_obj.FlowState.ERROR
    flow_obj.error_message = message
    flow_obj.backtrace = backtrace
    data_store.REL_DB.UpdateFlow(client_id, flow_id, flow_obj=flow_obj)

  def GetHuntResults(self, hunt_id):
    """Gets flow results for a given flow.

    Args:
      hunt_id: String with a hunt id or a RDFURN.

    Returns:
      List with hunt results payloads.
    """
    return data_store.REL_DB.ReadHuntResults(hunt_id, 0, sys.maxsize)


class DummyHuntOutputPlugin(output_plugin.OutputPlugin):
  num_calls = 0
  num_responses = 0

  def ProcessResponses(self, state, responses):
    DummyHuntOutputPlugin.num_calls += 1
    DummyHuntOutputPlugin.num_responses += len(list(responses))


class FailingDummyHuntOutputPlugin(output_plugin.OutputPlugin):

  def ProcessResponses(self, state, responses):
    raise RuntimeError("Oh no!")


class FailingInFlushDummyHuntOutputPlugin(output_plugin.OutputPlugin):

  def ProcessResponses(self, state, responses):
    pass

  def Flush(self, state):
    raise RuntimeError("Flush, oh no!")


class StatefulDummyHuntOutputPlugin(output_plugin.OutputPlugin):
  """Stateful dummy hunt output plugin."""
  data = []

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.delta = 0

  def InitializeState(self, state):
    super().InitializeState(state)
    state.index = 0

  def ProcessResponses(self, state, responses):
    StatefulDummyHuntOutputPlugin.data.append(state.index + self.delta)
    self.delta += 1

  def UpdateState(self, state):
    state.index += self.delta


class LongRunningDummyHuntOutputPlugin(output_plugin.OutputPlugin):
  num_calls = 0
  faketime = None

  def ProcessResponses(self, state, responses):
    LongRunningDummyHuntOutputPlugin.num_calls += 1
    LongRunningDummyHuntOutputPlugin.faketime.time += 100
