#!/usr/bin/env python
"""Tests for the standard hunts."""



import glob
import logging
import math
import os
import time


import mock

from grr.lib import flags
from grr.lib import rdfvalue
from grr.lib import stats
from grr.lib import utils
from grr.lib.rdfvalues import client as rdf_client
from grr.lib.rdfvalues import file_finder as rdf_file_finder
from grr.lib.rdfvalues import flows as rdf_flows
from grr.lib.rdfvalues import paths as rdf_paths
from grr.server import access_control
from grr.server import aff4
from grr.server import flow
from grr.server import foreman as rdf_foreman
from grr.server import output_plugin
from grr.server import queue_manager
from grr.server.aff4_objects import aff4_grr
from grr.server.aff4_objects import users as aff4_users
from grr.server.flows.general import administrative
from grr.server.flows.general import file_finder
from grr.server.flows.general import processes
from grr.server.flows.general import processes_test
from grr.server.flows.general import transfer
from grr.server.hunts import implementation
from grr.server.hunts import process_results
from grr.server.hunts import standard
from grr.test_lib import acl_test_lib
from grr.test_lib import action_mocks
from grr.test_lib import flow_test_lib
from grr.test_lib import hunt_test_lib
from grr.test_lib import test_lib


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

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


class FailingDummyHuntOutputPlugin(output_plugin.OutputPlugin):

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


class FailingInFlushDummyHuntOutputPlugin(output_plugin.OutputPlugin):

  def ProcessResponses(self, unused_responses):
    pass

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


class StatefulDummyHuntOutputPlugin(output_plugin.OutputPlugin):
  data = []

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

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


class LongRunningDummyHuntOutputPlugin(output_plugin.OutputPlugin):
  num_calls = 0

  def ProcessResponses(self, unused_responses):
    LongRunningDummyHuntOutputPlugin.num_calls += 1
    time.time = lambda: 100


class VerifiableDummyHuntOutputPlugin(output_plugin.OutputPlugin):

  def ProcessResponses(self, unused_responses):
    pass


class VerifiableDummyHuntOutputPluginVerfier(
    output_plugin.OutputPluginVerifier):
  plugin_name = VerifiableDummyHuntOutputPlugin.__name__

  num_calls = 0

  def VerifyHuntOutput(self, plugin, hunt):
    # Check that we get the plugin object we expected to get.
    # Actual verifiers implementations don't have to do this check.
    if not isinstance(plugin, VerifiableDummyHuntOutputPlugin):
      raise ValueError(
          "Passed plugin must be an "
          "VerifiableDummyHuntOutputPlugin, got: " % plugin.__class__.__name__)

    VerifiableDummyHuntOutputPluginVerfier.num_calls += 1
    return output_plugin.OutputPluginVerificationResult(
        status="SUCCESS", status_message="yo")

  def VerifyFlowOutput(self, plugin, hunt):
    pass


class DummyHuntOutputPluginWithRaisingVerifier(output_plugin.OutputPlugin):

  def ProcessResponses(self, unused_responses):
    pass


class DummyHuntOutputPluginWithRaisingVerifierVerifier(
    output_plugin.OutputPluginVerifier):
  plugin_name = DummyHuntOutputPluginWithRaisingVerifier.__name__

  def VerifyHuntOutput(self, plugin, hunt):
    # Check that we get the plugin object we expected to get.
    # Actual verifiers implementations don't have to do this check.
    if not isinstance(plugin, DummyHuntOutputPluginWithRaisingVerifier):
      raise ValueError(
          "Passed plugin must be an "
          "VerifiableDummyHuntOutputPlugin, got: " % plugin.__class__.__name__)

    raise RuntimeError("foobar")

  def VerifyFlowOutput(self, plugin, hunt):
    pass


class InfiniteFlow(flow.GRRFlow):
  """Flow that never ends."""

  @flow.StateHandler()
  def Start(self):
    self.CallState(next_state="NextState")

  @flow.StateHandler()
  def NextState(self):
    self.CallState(next_state="Start")


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

  def CreateHunt(self,
                 flow_runner_args=None,
                 flow_args=None,
                 client_rule_set=None,
                 original_object=None,
                 token=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_flows.FlowRunnerArgs(flow_name=transfer.GetFile.__name__))

    client_rule_set = (
        client_rule_set or
        rdf_foreman.ForemanClientRuleSet(rules=[
            rdf_foreman.ForemanClientRule(
                rule_type=rdf_foreman.ForemanClientRule.Type.REGEX,
                regex=rdf_foreman.ForemanRegexClientRule(
                    attribute_name="GRR client", attribute_regex="GRR"))
        ]))
    return implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=flow_runner_args,
        flow_args=flow_args,
        client_rule_set=client_rule_set,
        client_rate=0,
        original_object=original_object,
        token=token or self.token,
        **kwargs)

  def StartHunt(self, **kwargs):
    with self.CreateHunt(**kwargs) as hunt:
      hunt.Run()

    return hunt.urn

  def AssignTasksToClients(self, client_ids=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 = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token)
    for client_id in client_ids:
      foreman.AssignTasksToClient(client_id)

  def RunHunt(self, client_ids=None, iteration_limit=None, **mock_kwargs):
    client_mock = hunt_test_lib.SampleHuntMock(**mock_kwargs)
    hunt_test_lib.TestHuntHelper(
        client_mock,
        client_ids or self.client_ids,
        check_flow_errors=False,
        iteration_limit=iteration_limit,
        token=self.token)

  def StopHunt(self, hunt_urn):
    # Stop the hunt now.
    with aff4.FACTORY.Open(
        hunt_urn, age=aff4.ALL_TIMES, mode="rw", token=self.token) as hunt_obj:
      hunt_obj.Stop()

  def ProcessHuntOutputPlugins(self, **flow_args):
    flow_urn = flow.GRRFlow.StartFlow(
        flow_name=process_results.ProcessHuntResultCollectionsCronFlow.__name__,
        token=self.token,
        **flow_args)
    for _ in flow_test_lib.TestFlowHelper(flow_urn, token=self.token):
      pass
    return flow_urn


class StandardHuntTest(flow_test_lib.FlowTestsBaseclass, StandardHuntTestMixin):
  """Tests the Hunt."""

  def setUp(self):
    super(StandardHuntTest, self).setUp()
    # Set up 10 clients.
    self.client_ids = self.SetupClients(10)

    DummyHuntOutputPlugin.num_calls = 0
    DummyHuntOutputPlugin.num_responses = 0
    StatefulDummyHuntOutputPlugin.data = []
    LongRunningDummyHuntOutputPlugin.num_calls = 0

    with test_lib.FakeTime(0):
      # Clean up the foreman to remove any rules.
      with aff4.FACTORY.Open(
          "aff4:/foreman", mode="rw", token=self.token) as foreman:
        foreman.Set(foreman.Schema.RULES())

    self.old_logging_error = logging.error
    logging.error = self.AssertNoCollectionCorruption

  def tearDown(self):
    super(StandardHuntTest, self).tearDown()

    logging.error = self.old_logging_error

  def AssertNoCollectionCorruption(self, message, *args, **kwargs):
    self.assertFalse(
        "Results collection was changed outside of hunt" in message)
    self.old_logging_error(message, *args, **kwargs)

  def testResultCounting(self):
    path = os.path.join(self.base_path, "hello*")
    num_files = len(glob.glob(path))
    self.assertGreater(num_files, 0)

    hunt = implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flows.FlowRunnerArgs(
            flow_name=file_finder.FileFinder.__name__),
        flow_args=rdf_file_finder.FileFinderArgs(
            paths=[path],
            action=rdf_file_finder.FileFinderAction(action_type="STAT"),
        ),
        client_rate=0,
        token=self.token)
    hunt.Run()

    client_ids = self.SetupClients(5)
    self.AssignTasksToClients(client_ids=client_ids)
    action_mock = action_mocks.FileFinderClientMock()
    hunt_test_lib.TestHuntHelper(action_mock, client_ids, token=self.token)

    hunt = aff4.FACTORY.Open(hunt.urn, token=self.token)
    self.assertEqual(hunt.context.clients_with_results_count, 5)
    self.assertEqual(hunt.context.results_count, 5 * num_files)

  def testCreatesSymlinksOnClientsForEveryStartedFlow(self):
    hunt_urn = self.StartHunt()
    self.AssignTasksToClients()
    self.RunHunt()

    for client_id in self.client_ids:
      flows_fd = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token)
      flows_urns = list(flows_fd.ListChildren())
      self.assertEqual(len(flows_urns), 1)
      self.assertEqual(flows_urns[0].Basename(), hunt_urn.Basename() + ":hunt")

      # Check that the object is a symlink.
      fd = aff4.FACTORY.Open(
          flows_urns[0], follow_symlinks=False, token=self.token)
      self.assertEqual(fd.Get(fd.Schema.TYPE), "AFF4Symlink")

      target = fd.Get(fd.Schema.SYMLINK_TARGET)
      # Check that the symlink points into the hunt's namespace.
      self.assertTrue(str(target).startswith(str(hunt_urn)))

  def testDeletesSymlinksOnClientsWhenGetsDeletedItself(self):
    hunt_urn = self.StartHunt()
    self.AssignTasksToClients()
    self.RunHunt()

    for client_id in self.client_ids:
      # Check that symlinks to hunt-initiated flows are there as the hunt is
      # not deleted yet.
      flows_fd = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token)
      self.assertTrue(list(flows_fd.ListChildren()))

    aff4.FACTORY.Delete(hunt_urn, token=self.token)

    for client_id in self.client_ids:
      # Check that symlinks to hunt-initiated flows were deleted.
      flows_fd = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token)
      self.assertFalse(list(flows_fd.ListChildren()))

  def testStoppingHuntMarksFlowsForTerminationAndCleansQueues(self):
    with implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flows.FlowRunnerArgs(
            flow_name=InfiniteFlow.__name__),
        client_rule_set=rdf_foreman.ForemanClientRuleSet(
            rules=[
                rdf_foreman.ForemanClientRule(
                    rule_type=rdf_foreman.ForemanClientRule.Type.REGEX,
                    regex=rdf_foreman.ForemanRegexClientRule(
                        attribute_name="GRR client", attribute_regex="GRR"))
            ]),
        client_rate=0,
        token=self.token) as hunt:
      hunt.Run()

    self.AssignTasksToClients()

    # Run long enough for InfiniteFlows to start.
    self.RunHunt(iteration_limit=len(self.client_ids) * 2)
    self.StopHunt(hunt.urn)

    # All flows states should be destroyed by now.
    # If something is wrong with the GenericHunt.Stop implementation,
    # this will run forever.
    self.RunHunt()

    for client_id in self.client_ids:
      flows_root = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token)
      flows_list = list(flows_root.ListChildren())
      # Only one flow (issued by the hunt) is expected.
      self.assertEqual(len(flows_list), 1)

      # Check that flow's queues are deleted.
      with queue_manager.QueueManager(token=self.token) as manager:
        req_resp = list(manager.FetchRequestsAndResponses(flows_list[0]))
        self.assertFalse(req_resp)

      flow_obj = aff4.FACTORY.Open(
          flows_list[0], aff4_type=InfiniteFlow, token=self.token)
      self.assertEqual(
          flow_obj.Get(flow_obj.Schema.PENDING_TERMINATION).reason,
          "Parent hunt stopped.")

  def testGenericHuntWithoutOutputPlugins(self):
    """This tests running the hunt on some clients."""
    hunt_urn = self.StartHunt()
    self.AssignTasksToClients()
    self.RunHunt()
    self.StopHunt(hunt_urn)
    self.ProcessHuntOutputPlugins()

    with aff4.FACTORY.Open(
        hunt_urn, age=aff4.ALL_TIMES, token=self.token) as hunt_obj:

      started, finished, errors = hunt_obj.GetClientsCounts()
      self.assertEqual(started, 10)
      self.assertEqual(finished, 10)
      self.assertEqual(errors, 5)

      # Results collection is always written, even if there are no output
      # plugins.
      collection = implementation.GRRHunt.ResultCollectionForHID(hunt_urn)

      # We should receive stat entries.
      i = 0
      for i, x in enumerate(collection):
        self.assertEqual(x.payload.__class__, rdf_client.StatEntry)
        self.assertEqual(
            x.payload.AFF4Path(x.source).Split(2)[-1], "fs/os/tmp/evil.txt")

      self.assertEqual(i, 4)

      per_type_collection = implementation.GRRHunt.TypedResultCollectionForHID(
          hunt_urn)

      for i, x in enumerate(per_type_collection):
        self.assertEqual(x.payload.__class__, rdf_client.StatEntry)
        self.assertEqual(
            x.payload.AFF4Path(x.source).Split(2)[-1], "fs/os/tmp/evil.txt")

      self.assertListEqual(
          list(per_type_collection.ListStoredTypes()),
          [rdf_client.StatEntry.__name__])

      self.assertEqual(hunt_obj.context.clients_with_results_count, 5)
      self.assertEqual(hunt_obj.context.results_count, 5)

  def testHuntWithoutForemanRules(self):
    """Check no foreman rules are created if we pass add_foreman_rules=False."""
    hunt_urn = self.StartHunt(add_foreman_rules=False)
    with aff4.FACTORY.Open(
        "aff4:/foreman",
        mode="r",
        token=self.token,
        aff4_type=aff4_grr.GRRForeman) as foreman:
      foreman_rules = foreman.Get(
          foreman.Schema.RULES, default=foreman.Schema.RULES())
      self.assertFalse(foreman_rules)

    self.AssignTasksToClients()
    self.RunHunt()
    self.StopHunt(hunt_urn)

    with aff4.FACTORY.Open(
        hunt_urn, age=aff4.ALL_TIMES, token=self.token) as hunt_obj:

      started, finished, errors = hunt_obj.GetClientsCounts()
      self.assertEqual(started, 0)
      self.assertEqual(finished, 0)
      self.assertEqual(errors, 0)

  def testProcessHunResultsCronFlowDoesNothingWhenThereAreNoResults(self):
    # There's no hunt, nothing. Just assert that cron job completes
    # successfully.
    self.ProcessHuntOutputPlugins()

  def testOutputPluginsProcessOnlyNewResultsOnEveryRun(self):
    hunt_urn = self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name="DummyHuntOutputPlugin")
    ])

    # Process hunt results.
    self.ProcessHuntOutputPlugins()

    # Check that nothing has happened because hunt hasn't reported any
    # results yet.
    self.assertEqual(DummyHuntOutputPlugin.num_calls, 0)
    self.assertEqual(DummyHuntOutputPlugin.num_responses, 0)

    # Process first 5 clients
    self.AssignTasksToClients(self.client_ids[:5])

    # Run the hunt.
    self.RunHunt(failrate=-1)

    # Although we call ProcessHuntResultCollectionsCronFlow multiple times, it
    # should only call actual plugin once.
    for _ in range(5):
      self.ProcessHuntOutputPlugins()

    self.assertEqual(DummyHuntOutputPlugin.num_calls, 1)
    self.assertEqual(DummyHuntOutputPlugin.num_responses, 5)

    # Process last 5 clients
    self.AssignTasksToClients(self.client_ids[5:])

    # Run the hunt.
    self.RunHunt(failrate=-1)

    # Although we call ProcessHuntResultCollectionsCronFlow multiple times, it
    # should only call actual plugin once.
    for _ in range(5):
      self.ProcessHuntOutputPlugins()

    self.assertEqual(DummyHuntOutputPlugin.num_calls, 2)
    self.assertEqual(DummyHuntOutputPlugin.num_responses, 10)
    hunt = aff4.FACTORY.Open(hunt_urn, token=self.token)
    self.assertEqual(hunt.context.clients_with_results_count, 10)
    self.assertEqual(hunt.context.results_count, 10)

  def testOutputPluginsProcessingStatusIsWrittenToStatusCollection(self):
    plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="DummyHuntOutputPlugin")
    hunt_urn = self.StartHunt(output_plugins=[plugin_descriptor])

    # Run the hunt and process output plugins.
    self.AssignTasksToClients(self.client_ids)
    self.RunHunt(failrate=-1)
    self.ProcessHuntOutputPlugins()

    status_collection = implementation.GRRHunt.PluginStatusCollectionForHID(
        hunt_urn)
    errors_collection = implementation.GRRHunt.PluginErrorCollectionForHID(
        hunt_urn)

    self.assertEqual(len(errors_collection), 0)
    self.assertEqual(len(status_collection), 1)

    self.assertEqual(status_collection[0].status, "SUCCESS")
    self.assertEqual(status_collection[0].batch_index, 0)
    self.assertEqual(status_collection[0].batch_size, 10)
    self.assertEqual(status_collection[0].plugin_descriptor, plugin_descriptor)

  def testMultipleOutputPluginsProcessingStatusAreWrittenToStatusCollection(
      self):
    plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="DummyHuntOutputPlugin")
    hunt_urn = self.StartHunt(output_plugins=[plugin_descriptor])

    # Run the hunt on first 4 clients and process output plugins.
    self.AssignTasksToClients(self.client_ids[:4])
    self.RunHunt(failrate=-1)
    self.ProcessHuntOutputPlugins()

    # Run the hunt on last 6 clients and process output plugins.
    self.AssignTasksToClients(self.client_ids[4:])
    self.RunHunt(failrate=-1)
    self.ProcessHuntOutputPlugins()

    status_collection = implementation.GRRHunt.PluginStatusCollectionForHID(
        hunt_urn)
    errors_collection = implementation.GRRHunt.PluginErrorCollectionForHID(
        hunt_urn)

    self.assertEqual(len(errors_collection), 0)
    self.assertEqual(len(status_collection), 2)

    items = sorted(status_collection, key=lambda x: x.age)
    self.assertEqual(items[0].status, "SUCCESS")
    self.assertEqual(items[0].batch_index, 0)
    self.assertEqual(items[0].batch_size, 4)
    self.assertEqual(items[0].plugin_descriptor, plugin_descriptor)

    self.assertEqual(items[1].status, "SUCCESS")
    self.assertEqual(items[1].batch_index, 0)
    self.assertEqual(items[1].batch_size, 6)
    self.assertEqual(items[1].plugin_descriptor, plugin_descriptor)

  def testErrorOutputPluginStatusIsAlsoWrittenToErrorsCollection(self):
    failing_plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="FailingDummyHuntOutputPlugin")
    plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="DummyHuntOutputPlugin")
    hunt_urn = self.StartHunt(
        output_plugins=[failing_plugin_descriptor, plugin_descriptor])

    # Run the hunt and process output plugins.
    self.AssignTasksToClients(self.client_ids)
    self.RunHunt(failrate=-1)
    try:
      self.ProcessHuntOutputPlugins()
    except process_results.ResultsProcessingError:
      pass

    status_collection = implementation.GRRHunt.PluginStatusCollectionForHID(
        hunt_urn)
    errors_collection = implementation.GRRHunt.PluginErrorCollectionForHID(
        hunt_urn)

    self.assertEqual(len(errors_collection), 1)
    self.assertEqual(len(status_collection), 2)

    self.assertEqual(errors_collection[0].status, "ERROR")
    self.assertEqual(errors_collection[0].batch_index, 0)
    self.assertEqual(errors_collection[0].batch_size, 10)
    self.assertEqual(errors_collection[0].plugin_descriptor,
                     failing_plugin_descriptor)
    self.assertEqual(errors_collection[0].summary, "Oh no!")

    items = sorted(
        status_collection, key=lambda x: x.plugin_descriptor.plugin_name)
    self.assertEqual(items[0].status, "SUCCESS")
    self.assertEqual(items[0].batch_index, 0)
    self.assertEqual(items[0].batch_size, 10)
    self.assertEqual(items[0].plugin_descriptor, plugin_descriptor)

    self.assertEqual(items[1].status, "ERROR")
    self.assertEqual(items[1].batch_index, 0)
    self.assertEqual(items[1].batch_size, 10)
    self.assertEqual(items[1].plugin_descriptor, failing_plugin_descriptor)
    self.assertEqual(items[1].summary, "Oh no!")

  def testOutputPluginFlushErrorIsLoggedProperly(self):
    failing_plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="FailingInFlushDummyHuntOutputPlugin")
    hunt_urn = self.StartHunt(output_plugins=[failing_plugin_descriptor])

    # Run the hunt and process output plugins.
    self.AssignTasksToClients(self.client_ids)
    self.RunHunt(failrate=-1)
    try:
      self.ProcessHuntOutputPlugins()
    except process_results.ResultsProcessingError:
      pass

    status_collection = implementation.GRRHunt.PluginStatusCollectionForHID(
        hunt_urn)
    errors_collection = implementation.GRRHunt.PluginErrorCollectionForHID(
        hunt_urn)

    self.assertEqual(len(errors_collection), 1)
    self.assertEqual(len(status_collection), 1)

    self.assertEqual(errors_collection[0].status, "ERROR")
    self.assertEqual(errors_collection[0].batch_index, 0)
    self.assertEqual(errors_collection[0].batch_size, 10)
    self.assertEqual(errors_collection[0].plugin_descriptor,
                     failing_plugin_descriptor)
    self.assertEqual(errors_collection[0].summary, "Flush, oh no!")

    items = sorted(
        status_collection, key=lambda x: x.plugin_descriptor.plugin_name)

    self.assertEqual(items[0].status, "ERROR")
    self.assertEqual(items[0].batch_index, 0)
    self.assertEqual(items[0].batch_size, 10)
    self.assertEqual(items[0].plugin_descriptor, failing_plugin_descriptor)
    self.assertEqual(items[0].summary, "Flush, oh no!")

  def testFailingOutputPluginDoesNotAffectOtherOutputPlugins(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name="FailingDummyHuntOutputPlugin"),
        output_plugin.OutputPluginDescriptor(
            plugin_name="DummyHuntOutputPlugin")
    ])

    # Process hunt results.
    self.ProcessHuntOutputPlugins()

    self.assertEqual(DummyHuntOutputPlugin.num_calls, 0)
    self.assertEqual(DummyHuntOutputPlugin.num_responses, 0)

    self.AssignTasksToClients()
    self.RunHunt(failrate=-1)

    # We shouldn't get any more calls after the first call to
    # ProcessHuntResultCollectionsCronFlow.
    self.assertRaises(process_results.ResultsProcessingError,
                      self.ProcessHuntOutputPlugins)
    for _ in range(5):
      self.ProcessHuntOutputPlugins()

    self.assertEqual(DummyHuntOutputPlugin.num_calls, 1)
    self.assertEqual(DummyHuntOutputPlugin.num_responses, 10)

  def testResultsProcessingErrorContainsDetailedFailureData(self):
    failing_plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="FailingDummyHuntOutputPlugin")
    plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="DummyHuntOutputPlugin")
    hunt_urn = self.StartHunt(
        output_plugins=[failing_plugin_descriptor, plugin_descriptor])

    # Process hunt results.
    self.ProcessHuntOutputPlugins()

    self.assertEqual(DummyHuntOutputPlugin.num_calls, 0)
    self.assertEqual(DummyHuntOutputPlugin.num_responses, 0)

    self.AssignTasksToClients()
    self.RunHunt(failrate=-1)

    # We shouldn't get any more calls after the first call to
    # ProcessHuntResultCollectionsCronFlow.
    try:
      self.ProcessHuntOutputPlugins()

      # We shouldn't get here.
      self.fail()
    except process_results.ResultsProcessingError as e:
      self.assertEqual(len(e.exceptions_by_hunt), 1)
      self.assertTrue(hunt_urn in e.exceptions_by_hunt)
      self.assertEqual(len(e.exceptions_by_hunt[hunt_urn]), 1)
      self.assertTrue(
          failing_plugin_descriptor in e.exceptions_by_hunt[hunt_urn])
      self.assertEqual(
          len(e.exceptions_by_hunt[hunt_urn][failing_plugin_descriptor]), 1)
      self.assertEqual(
          e.exceptions_by_hunt[hunt_urn][failing_plugin_descriptor][0].message,
          "Oh no!")

  @mock.patch.object(
      FailingDummyHuntOutputPlugin,
      "ProcessResponses",
      side_effect=RuntimeError("Oh, no"))
  def testResultsAreNotProcessedAgainAfterPluginFailure(self,
                                                        process_responses_mock):
    failing_plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="FailingDummyHuntOutputPlugin")
    self.StartHunt(output_plugins=[failing_plugin_descriptor])

    self.AssignTasksToClients()
    self.RunHunt(failrate=-1)

    # Process hunt results.
    try:
      self.ProcessHuntOutputPlugins()
    except process_results.ResultsProcessingError:
      pass
    self.assertEqual(process_responses_mock.call_count, 1)

    self.ProcessHuntOutputPlugins()
    # Check that call count hasn't changed.
    self.assertEqual(process_responses_mock.call_count, 1)

  def testUpdatesStatsCounterOnSuccess(self):
    failing_plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="DummyHuntOutputPlugin")
    self.StartHunt(output_plugins=[failing_plugin_descriptor])

    prev_success_count = stats.STATS.GetMetricValue(
        "hunt_results_ran_through_plugin", fields=["DummyHuntOutputPlugin"])
    prev_errors_count = stats.STATS.GetMetricValue(
        "hunt_output_plugin_errors", fields=["DummyHuntOutputPlugin"])

    self.AssignTasksToClients()
    self.RunHunt(failrate=-1)
    self.ProcessHuntOutputPlugins()

    success_count = stats.STATS.GetMetricValue(
        "hunt_results_ran_through_plugin", fields=["DummyHuntOutputPlugin"])
    errors_count = stats.STATS.GetMetricValue(
        "hunt_output_plugin_errors", fields=["DummyHuntOutputPlugin"])

    # 1 result for each client makes it 10 results.
    self.assertEqual(success_count - prev_success_count, 10)

    self.assertEqual(errors_count - prev_errors_count, 0)

  def testUpdatesStatsCounterOnFailure(self):
    failing_plugin_descriptor = output_plugin.OutputPluginDescriptor(
        plugin_name="FailingDummyHuntOutputPlugin")
    self.StartHunt(output_plugins=[failing_plugin_descriptor])

    prev_success_count = stats.STATS.GetMetricValue(
        "hunt_results_ran_through_plugin",
        fields=["FailingDummyHuntOutputPlugin"])
    prev_errors_count = stats.STATS.GetMetricValue(
        "hunt_output_plugin_errors", fields=["FailingDummyHuntOutputPlugin"])

    self.AssignTasksToClients()
    self.RunHunt(failrate=-1)
    try:
      self.ProcessHuntOutputPlugins()
    except process_results.ResultsProcessingError:
      pass

    success_count = stats.STATS.GetMetricValue(
        "hunt_results_ran_through_plugin",
        fields=["FailingDummyHuntOutputPlugin"])
    errors_count = stats.STATS.GetMetricValue(
        "hunt_output_plugin_errors", fields=["FailingDummyHuntOutputPlugin"])

    self.assertEqual(success_count - prev_success_count, 0)
    self.assertEqual(errors_count - prev_errors_count, 1)

  def testOutputPluginsMaintainState(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name="StatefulDummyHuntOutputPlugin")
    ])

    self.assertListEqual(StatefulDummyHuntOutputPlugin.data, [])

    # Run the hunt on every client and separately and run the output
    # cron flow for every client to ensure that output plugin will
    # run multiple times.
    for index in range(10):
      self.AssignTasksToClients([self.client_ids[index]])

      # Run the hunt.
      self.RunHunt(failrate=-1)
      self.ProcessHuntOutputPlugins()

    # Output plugins should have been called 10 times, adding a number
    # to the "data" list on every call and incrementing it each time.
    self.assertListEqual(StatefulDummyHuntOutputPlugin.data,
                         [0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

  def testMultipleHuntsOutputIsProcessedCorrectly(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name="DummyHuntOutputPlugin")
    ])
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name="StatefulDummyHuntOutputPlugin")
    ])

    self.AssignTasksToClients()
    self.RunHunt(failrate=-1)
    self.ProcessHuntOutputPlugins()

    # Check that plugins worked correctly
    self.assertEqual(DummyHuntOutputPlugin.num_calls, 1)
    self.assertListEqual(StatefulDummyHuntOutputPlugin.data, [0])

  def testProcessHuntResultCollectionsCronFlowAbortsIfRunningTooLong(self):
    self.assertEqual(LongRunningDummyHuntOutputPlugin.num_calls, 0)

    test = [0]

    def TimeStub():
      test[0] += 1e-6
      return test[0]

    with utils.Stubber(time, "time", TimeStub):
      self.StartHunt(output_plugins=[
          output_plugin.OutputPluginDescriptor(
              plugin_name="LongRunningDummyHuntOutputPlugin")
      ])
      self.AssignTasksToClients()
      self.RunHunt(failrate=-1)

      # LongRunningDummyHuntOutputPlugin will set the time to 100s on the first
      # run, which will effectively mean that it's running for too long.
      self.ProcessHuntOutputPlugins(
          batch_size=1, max_running_time=rdfvalue.Duration("99s"))

      # In normal conditions, there should be 10 results generated.
      # With batch size of 1 this should result in 10 calls to output plugin.
      # But as we were using TimeStub, the flow should have aborted after 1
      # call.
      self.assertEqual(LongRunningDummyHuntOutputPlugin.num_calls, 1)

  def testProcessHuntResultCollectionsCronFlowDoesNotAbortsIfRunningInTime(
      self):
    self.assertEqual(LongRunningDummyHuntOutputPlugin.num_calls, 0)

    test = [0]

    def TimeStub():
      test[0] += 1e-6
      return test[0]

    with utils.Stubber(time, "time", TimeStub):
      self.StartHunt(output_plugins=[
          output_plugin.OutputPluginDescriptor(
              plugin_name="LongRunningDummyHuntOutputPlugin")
      ])
      self.AssignTasksToClients()
      self.RunHunt(failrate=-1)

      # LongRunningDummyHuntOutputPlugin will set the time to 100s on the first
      # run, which will effectively mean that it's running in time.
      self.ProcessHuntOutputPlugins(
          batch_size=1, max_running_time=rdfvalue.Duration("101s"))

      # In normal conditions, there should be 10 results generated.
      self.assertEqual(LongRunningDummyHuntOutputPlugin.num_calls, 10)

  def testHuntResultsArrivingWhileOldResultsAreProcessedAreHandled(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name="DummyHuntOutputPlugin")
    ])

    # Process hunt results.
    self.ProcessHuntOutputPlugins()

    # Check that nothing has happened because hunt hasn't reported any
    # results yet.
    self.assertEqual(DummyHuntOutputPlugin.num_calls, 0)
    self.assertEqual(DummyHuntOutputPlugin.num_responses, 0)

    # Generate new results while the plugin is working.
    self.num_processed = 0

    def ProcessResponsesStub(_, responses):
      # Add 5 more results the first time we are called.
      if not self.num_processed:
        self.AssignTasksToClients(self.client_ids[5:])
        self.RunHunt(failrate=-1)
      # Just count the total number processed - we don't care about batch size
      # at this point.
      self.num_processed += len(responses)

    with utils.Stubber(DummyHuntOutputPlugin, "ProcessResponses",
                       ProcessResponsesStub):
      self.AssignTasksToClients(self.client_ids[:5])
      self.RunHunt(failrate=-1)
      self.ProcessHuntOutputPlugins()

    self.assertEqual(10, self.num_processed)
    del self.num_processed

  def _AppendFlowRequest(self, flows, client_id, file_id):
    flows.Append(
        client_ids=["C.1%015d" % client_id],
        runner_args=rdf_flows.FlowRunnerArgs(
            flow_name=transfer.GetFile.__name__),
        args=transfer.GetFileArgs(
            pathspec=rdf_paths.PathSpec(
                path="/tmp/evil%s.txt" % file_id,
                pathtype=rdf_paths.PathSpec.PathType.OS),))

  def RunVariableGenericHunt(self):
    args = standard.VariableGenericHuntArgs()
    self._AppendFlowRequest(args.flows, 1, 1)
    self._AppendFlowRequest(args.flows, 2, 2)
    self._AppendFlowRequest(args.flows, 2, 3)

    with implementation.GRRHunt.StartHunt(
        hunt_name=standard.VariableGenericHunt.__name__,
        args=args,
        client_rate=0,
        token=self.token) as hunt:
      hunt.Run()
      hunt.ManuallyScheduleClients()

    # Run the hunt.
    client_mock = hunt_test_lib.SampleHuntMock(failrate=100)
    hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False,
                                 self.token)

    with aff4.FACTORY.Open(
        hunt.session_id, mode="rw", token=self.token) as hunt:
      hunt.Stop()

    return hunt

  def testVariableGenericHunt(self):
    """This tests running the hunt on some clients."""
    hunt = self.RunVariableGenericHunt()

    hunt_obj = aff4.FACTORY.Open(
        hunt.session_id, age=aff4.ALL_TIMES, token=self.token)
    started, finished, errors = hunt_obj.GetClientsCounts()
    self.assertEqual(started, 2)
    # Amazing as it may sound, 3 is actually a correct value as we run 2 flows
    # on a second client.
    self.assertEqual(finished, 3)
    self.assertEqual(errors, 0)

  def testHuntTermination(self):
    """This tests that hunts with a client limit terminate correctly."""
    with test_lib.FakeTime(1000, increment=1e-6):
      with implementation.GRRHunt.StartHunt(
          hunt_name=standard.GenericHunt.__name__,
          flow_runner_args=rdf_flows.FlowRunnerArgs(
              flow_name=transfer.GetFile.__name__),
          flow_args=transfer.GetFileArgs(pathspec=rdf_paths.PathSpec(
              path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)),
          client_rule_set=rdf_foreman.ForemanClientRuleSet(
              rules=[
                  rdf_foreman.ForemanClientRule(
                      rule_type=rdf_foreman.ForemanClientRule.Type.REGEX,
                      regex=rdf_foreman.ForemanRegexClientRule(
                          attribute_name="GRR client", attribute_regex="GRR"))
              ]),
          client_limit=5,
          client_rate=0,
          expiry_time=rdfvalue.Duration("1000s"),
          token=self.token) as hunt:
        hunt.Run()

      # Pretend to be the foreman now and dish out hunting jobs to all the
      # clients (Note we have 10 clients here).
      foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token)
      for client_id in self.client_ids:
        foreman.AssignTasksToClient(client_id)

      # Run the hunt.
      client_mock = hunt_test_lib.SampleHuntMock()
      hunt_test_lib.TestHuntHelper(
          client_mock,
          self.client_ids,
          check_flow_errors=False,
          token=self.token)

      hunt_obj = aff4.FACTORY.Open(
          hunt.session_id, age=aff4.ALL_TIMES, token=self.token)

      started, finished, errors = hunt_obj.GetClientsCounts()
      self.assertEqual(started, 5)
      self.assertEqual(finished, 5)
      self.assertEqual(errors, 2)

      hunt_obj = aff4.FACTORY.Open(
          hunt.session_id, age=aff4.ALL_TIMES, token=self.token)

      # Hunts are automatically paused when they reach the client limit.
      self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "PAUSED")

  def testHuntIsStoppedIfCrashNumberOverThreshold(self):
    with self.CreateHunt(crash_limit=3, token=self.token) as hunt:
      hunt.Run()

    # Run the hunt on 2 clients.
    for client_id in self.client_ids[:2]:
      self.AssignTasksToClients([client_id])
      client_mock = flow_test_lib.CrashClientMock(client_id, token=self.token)
      hunt_test_lib.TestHuntHelper(
          client_mock, [client_id], check_flow_errors=False, token=self.token)

    # Hunt should still be running: 2 crashes are within the threshold.
    hunt_obj = aff4.FACTORY.Open(
        hunt.session_id, age=aff4.ALL_TIMES, token=self.token)
    self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "STARTED")

    # Run the hunt on another client.
    client_id = self.client_ids[2]
    client_mock = flow_test_lib.CrashClientMock(client_id, token=self.token)
    self.AssignTasksToClients([client_id])
    hunt_test_lib.TestHuntHelper(
        client_mock, [client_id], check_flow_errors=False, token=self.token)

    # Hunt should be terminated: 3 crashes are over the threshold.
    hunt_obj = aff4.FACTORY.Open(
        hunt.session_id, age=aff4.ALL_TIMES, token=self.token)
    self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "STOPPED")

  def _CheckHuntStoppedNotification(self, str_match):
    user_record = aff4.FACTORY.Create(
        aff4.ROOT_URN.Add("users").Add(self.token.username),
        aff4_type=aff4_users.GRRUser,
        mode="r",
        token=self.token)
    pending = user_record.Get(user_record.Schema.PENDING_NOTIFICATIONS)
    self.assertEqual(len(pending), 1)
    self.assertTrue(str_match in pending[0].message)

  def testHuntIsStoppedIfAveragePerClientResultsCountTooHigh(self):
    with utils.Stubber(implementation.GRRHunt,
                       "MIN_CLIENTS_FOR_AVERAGE_THRESHOLDS", 4):

      flow_args = processes.ListProcessesArgs()
      flow_runner_args = rdf_flows.FlowRunnerArgs(
          flow_name=processes.ListProcesses.__name__)

      hunt_urn = self.StartHunt(
          flow_args=flow_args,
          flow_runner_args=flow_runner_args,
          avg_results_per_client_limit=1,
          token=self.token)

      def RunOnClients(client_ids, num_processes):
        client_mock = processes_test.ListProcessesMock(
            [rdf_client.Process(pid=1, exe="a.exe")] * num_processes)
        self.AssignTasksToClients(client_ids)
        hunt_test_lib.TestHuntHelper(
            client_mock, client_ids, check_flow_errors=False, token=self.token)

      def CheckState(expected_state, expected_results_count):
        hunt_obj = aff4.FACTORY.Open(hunt_urn, token=self.token)
        self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), expected_state)
        self.assertEqual(hunt_obj.context.results_count, expected_results_count)

      RunOnClients(self.client_ids[:2], 1)
      # Hunt should still be running: we got 1 response from 2 clients. We need
      # at least 3 clients to start calculating the average.
      CheckState("STARTED", 2)

      RunOnClients([self.client_ids[2]], 2)
      # Hunt should still be running: we got 1 response for first 2 clients and
      # 2 responses for the third. This is over the limit but we need at least 4
      # clients to start applying thresholds.
      CheckState("STARTED", 4)

      RunOnClients([self.client_ids[3]], 0)
      # Hunt should still be running: we got 1 response for first 2 clients,
      # 2 responses for the third and zero for the 4th. This makes it 1 result
      # per client on average. This is within the limit of 1.
      CheckState("STARTED", 4)

      RunOnClients(self.client_ids[4:5], 2)
      # Hunt should be terminated: 5 clients did run and we got 6 results.
      # That's more than the allowed average of 1.
      # Note that this check also implicitly checks that the 6th client didn't
      # run at all (otherwise total number of results would be 8, not 6).
      CheckState("STOPPED", 6)

      self._CheckHuntStoppedNotification(
          "reached the average results per client")

  def testHuntIsStoppedIfAveragePerClientCpuUsageTooHigh(self):
    with utils.Stubber(implementation.GRRHunt,
                       "MIN_CLIENTS_FOR_AVERAGE_THRESHOLDS", 4):
      hunt_urn = self.StartHunt(
          avg_cpu_seconds_per_client_limit=3, token=self.token)

      def RunOnClients(client_ids, user_cpu_time, system_cpu_time):
        self.AssignTasksToClients(client_ids)
        self.RunHunt(
            client_ids=client_ids,
            user_cpu_time=user_cpu_time,
            system_cpu_time=system_cpu_time)

      def CheckState(expected_state, expected_user_cpu, expected_system_cpu):
        hunt_obj = aff4.FACTORY.Open(hunt_urn, token=self.token)
        self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), expected_state)
        self.assertEqual(
            hunt_obj.context.client_resources.cpu_usage.user_cpu_time,
            expected_user_cpu)
        self.assertEqual(
            hunt_obj.context.client_resources.cpu_usage.system_cpu_time,
            expected_system_cpu)

      RunOnClients(self.client_ids[:2], 1, 2)
      # Hunt should still be running: we need at least 3 clients to start
      # calculating the average.
      CheckState("STARTED", 2, 4)

      RunOnClients([self.client_ids[2]], 2, 4)
      # Hunt should still be running: even though the average is higher than the
      # limit, number of clients is not enough.
      CheckState("STARTED", 4, 8)

      RunOnClients([self.client_ids[3]], 0, 0)
      # Hunt should still be running: we got 4 clients, which is enough to check
      # average per-client CPU usage. But 4 user cpu + 8 system cpu seconds for
      # 4 clients make an average of 3 seconds per client - this is within the
      # limit.
      CheckState("STARTED", 4, 8)

      RunOnClients([self.client_ids[4]], 2, 4)
      # Hunt should be terminated: the average is exceeded.
      CheckState("STOPPED", 6, 12)

      self._CheckHuntStoppedNotification(
          "reached the average CPU seconds per client")

  def testHuntIsStoppedIfAveragePerClientNetworkUsageTooHigh(self):
    with utils.Stubber(implementation.GRRHunt,
                       "MIN_CLIENTS_FOR_AVERAGE_THRESHOLDS", 4):
      hunt_urn = self.StartHunt(
          avg_network_bytes_per_client_limit=1, token=self.token)

      def RunOnClients(client_ids, network_bytes_sent):
        self.AssignTasksToClients(client_ids)
        self.RunHunt(
            client_ids=client_ids, network_bytes_sent=network_bytes_sent)

      def CheckState(expected_state, expected_network_bytes_sent):
        hunt_obj = aff4.FACTORY.Open(hunt_urn, token=self.token)
        self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), expected_state)
        self.assertEqual(hunt_obj.context.network_bytes_sent,
                         expected_network_bytes_sent)

      RunOnClients(self.client_ids[:2], 1)
      # Hunt should still be running: we need at least 3 clients to start
      # calculating the average.
      CheckState("STARTED", 2)

      RunOnClients([self.client_ids[2]], 2)
      # Hunt should still be running: even though the average is higher than the
      # limit, number of clients is not enough.
      CheckState("STARTED", 4)

      RunOnClients([self.client_ids[3]], 0)
      # Hunt should still be running: we got 4 clients, which is enough to check
      # average per-client network bytes usage, but 4 bytes for 4 clients is
      # within the limit of 1 byte per client on average.
      CheckState("STARTED", 4)

      RunOnClients([self.client_ids[4]], 2)
      # Hunt should be terminated: the limit is exceeded.
      CheckState("STOPPED", 6)

      self._CheckHuntStoppedNotification(
          "reached the average network bytes per client")

  def testHuntExpiration(self):
    """This tests that hunts with a client limit terminate correctly."""
    with test_lib.FakeTime(1000):
      with implementation.GRRHunt.StartHunt(
          hunt_name=standard.GenericHunt.__name__,
          flow_runner_args=rdf_flows.FlowRunnerArgs(
              flow_name=transfer.GetFile.__name__),
          flow_args=transfer.GetFileArgs(pathspec=rdf_paths.PathSpec(
              path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)),
          client_rule_set=rdf_foreman.ForemanClientRuleSet(
              rules=[
                  rdf_foreman.ForemanClientRule(
                      rule_type=rdf_foreman.ForemanClientRule.Type.REGEX,
                      regex=rdf_foreman.ForemanRegexClientRule(
                          attribute_name="GRR client", attribute_regex="GRR"))
              ]),
          client_limit=5,
          expiry_time=rdfvalue.Duration("1000s"),
          token=self.token) as hunt:
        hunt.Run()

      # Pretend to be the foreman now and dish out hunting jobs to all the
      # clients (Note we have 10 clients here).
      foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token)
      for client_id in self.client_ids:
        foreman.AssignTasksToClient(client_id)

      hunt_obj = aff4.FACTORY.Open(
          hunt.session_id, age=aff4.ALL_TIMES, token=self.token)

      self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "STARTED")

      # Now advance the time such that the hunt expires.
      time.time = lambda: 5000

      # Run the hunt.
      client_mock = hunt_test_lib.SampleHuntMock()
      hunt_test_lib.TestHuntHelper(
          client_mock,
          self.client_ids,
          check_flow_errors=False,
          token=self.token)

      # No client should be processed since the hunt is expired.
      started, finished, errors = hunt_obj.GetClientsCounts()
      self.assertEqual(started, 0)
      self.assertEqual(finished, 0)
      self.assertEqual(errors, 0)

      hunt_obj = aff4.FACTORY.Open(
          hunt.session_id, age=aff4.ALL_TIMES, token=self.token)

      # Hunts are automatically stopped when they expire.
      self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "COMPLETED")

  def testHuntModificationWorksCorrectly(self):
    """This tests running the hunt on some clients."""
    with implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flows.FlowRunnerArgs(
            flow_name=transfer.GetFile.__name__),
        flow_args=transfer.GetFileArgs(
            pathspec=rdf_paths.PathSpec(
                path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS),
        ),
        client_rule_set=rdf_foreman.ForemanClientRuleSet(
            rules=[
                rdf_foreman.ForemanClientRule(
                    rule_type=rdf_foreman.ForemanClientRule.Type.REGEX,
                    regex=rdf_foreman.ForemanRegexClientRule(
                        attribute_name="GRR client", attribute_regex="GRR"))
            ]),
        client_limit=1,
        client_rate=0,
        token=self.token) as hunt:
      hunt.Run()

    # Forget about hunt object, we'll use AFF4 for everything.
    hunt_session_id = hunt.session_id
    hunt = None

    # Pretend to be the foreman now and dish out hunting jobs to all the
    # client..
    with aff4.FACTORY.Open(
        "aff4:/foreman", mode="rw", token=self.token) as foreman:
      for client_id in self.client_ids:
        foreman.AssignTasksToClient(client_id)

    # Run the hunt.
    client_mock = hunt_test_lib.SampleHuntMock()
    hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False,
                                 self.token)

    # Re-open the hunt to get fresh data.
    hunt_obj = aff4.FACTORY.Open(
        hunt_session_id, age=aff4.ALL_TIMES, token=self.token)

    # There should be only one client, due to the limit
    started, _, _ = hunt_obj.GetClientsCounts()
    self.assertEqual(started, 1)

    # Check the hunt is paused.
    self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "PAUSED")

    with aff4.FACTORY.Open(
        hunt_session_id, mode="rw", token=self.token) as hunt_obj:
      runner = hunt_obj.GetRunner()
      runner.runner_args.client_limit = 10
      runner.Start()

    # Pretend to be the foreman now and dish out hunting jobs to all the
    # clients.
    with aff4.FACTORY.Open(
        "aff4:/foreman", mode="rw", token=self.token) as foreman:
      for client_id in self.client_ids:
        foreman.AssignTasksToClient(client_id)

    hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False,
                                 self.token)

    hunt_obj = aff4.FACTORY.Open(
        hunt_session_id, age=aff4.ALL_TIMES, token=self.token)
    # There should be only one client, due to the limit
    started, _, _ = hunt_obj.GetClientsCounts()
    self.assertEqual(started, 10)

  def testResourceUsageStats(self):
    client_ids = self.SetupClients(10)

    with implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flows.FlowRunnerArgs(
            flow_name=transfer.GetFile.__name__),
        flow_args=transfer.GetFileArgs(pathspec=rdf_paths.PathSpec(
            path="/tmp/evil.txt",
            pathtype=rdf_paths.PathSpec.PathType.OS,
        )),
        client_rule_set=rdf_foreman.ForemanClientRuleSet(
            rules=[
                rdf_foreman.ForemanClientRule(
                    rule_type=rdf_foreman.ForemanClientRule.Type.REGEX,
                    regex=rdf_foreman.ForemanRegexClientRule(
                        attribute_name="GRR client", attribute_regex="GRR"))
            ]),
        output_plugins=[],
        client_rate=0,
        token=self.token) as hunt:
      hunt.Run()

    with aff4.FACTORY.Open(
        "aff4:/foreman", mode="rw", token=self.token) as foreman:
      for client_id in client_ids:
        foreman.AssignTasksToClient(client_id)

    client_mock = hunt_test_lib.SampleHuntMock()
    hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, self.token)

    hunt = aff4.FACTORY.Open(
        hunt.urn, aff4_type=standard.GenericHunt, token=self.token)

    # This is called once for each state method. Each flow above runs the
    # Start and the StoreResults methods.
    usage_stats = hunt.context.usage_stats
    self.assertEqual(usage_stats.user_cpu_stats.num, 10)
    self.assertTrue(math.fabs(usage_stats.user_cpu_stats.mean - 5.5) < 1e-7)
    self.assertTrue(
        math.fabs(usage_stats.user_cpu_stats.std - 2.8722813) < 1e-7)

    self.assertEqual(usage_stats.system_cpu_stats.num, 10)
    self.assertTrue(math.fabs(usage_stats.system_cpu_stats.mean - 11) < 1e-7)
    self.assertTrue(
        math.fabs(usage_stats.system_cpu_stats.std - 5.7445626) < 1e-7)

    self.assertEqual(usage_stats.network_bytes_sent_stats.num, 10)
    self.assertTrue(
        math.fabs(usage_stats.network_bytes_sent_stats.mean - 16.5) < 1e-7)
    self.assertTrue(
        math.fabs(usage_stats.network_bytes_sent_stats.std - 8.61684396) < 1e-7)

    # NOTE: Not checking histograms here. RunningStatsTest tests that mean,
    # standard deviation and histograms are calculated correctly. Therefore
    # if mean/stdev values are correct histograms should be ok as well.

    self.assertEqual(len(usage_stats.worst_performers), 10)

    prev = usage_stats.worst_performers[0]
    for p in usage_stats.worst_performers[1:]:
      self.assertTrue(
          prev.cpu_usage.user_cpu_time + prev.cpu_usage.system_cpu_time >
          p.cpu_usage.user_cpu_time + p.cpu_usage.system_cpu_time)
      prev = p

  def testHuntCollectionLogging(self):
    """This tests running the hunt on some clients."""
    with implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flows.FlowRunnerArgs(
            flow_name=flow_test_lib.DummyLogFlow.__name__),
        client_rate=0,
        token=self.token) as hunt:
      hunt.Run()
      hunt.Log("Log from the hunt itself")

    hunt_urn = hunt.urn

    self.AssignTasksToClients()
    self.RunHunt()

    # Check logs were written to the hunt collection
    hunt_logs = implementation.GRRHunt.LogCollectionForHID(hunt_urn)
    count = 0
    for log in hunt_logs:
      if log.client_id:
        self.assertTrue(log.client_id in self.client_ids)
        self.assertTrue(log.log_message in [
            "First", "Second", "Third", "Fourth", "Uno", "Dos", "Tres", "Cuatro"
        ])
        self.assertTrue(log.flow_name in [
            flow_test_lib.DummyLogFlow.__name__,
            flow_test_lib.DummyLogFlowChild.__name__
        ])
        self.assertTrue(str(hunt_urn) in str(log.urn))
      else:
        self.assertEqual(log.log_message, "Log from the hunt itself")
        self.assertEqual(log.flow_name, standard.GenericHunt.__name__)
        self.assertEqual(log.urn, hunt_urn)

      count += 1
    # 4 logs for each flow, 2 flow run.  One hunt-level log.
    self.assertEqual(count, 8 * len(self.client_ids) + 1)

  def testCreatorPropagation(self):
    self.CreateAdminUser("adminuser")
    admin_token = access_control.ACLToken(
        username="adminuser", reason="testing")
    # Start a flow that requires admin privileges in the hunt. The
    # parameters are not valid so the flow will error out but it's
    # enough to check if the flow was actually run (i.e., it passed
    # the label test).
    with implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flows.FlowRunnerArgs(
            flow_name=administrative.UpdateClient.__name__),
        flow_args=administrative.UpdateClientArgs(),
        client_rule_set=rdf_foreman.ForemanClientRuleSet(
            rules=[
                rdf_foreman.ForemanClientRule(
                    rule_type=rdf_foreman.ForemanClientRule.Type.REGEX,
                    regex=rdf_foreman.ForemanRegexClientRule(
                        attribute_name="GRR client", attribute_regex="GRR"))
            ]),
        client_rate=0,
        token=admin_token) as hunt:
      hunt.Run()

    self.CreateUser("nonadmin")
    nonadmin_token = access_control.ACLToken(
        username="nonadmin", reason="testing")
    self.AssignTasksToClients()

    client_mock = hunt_test_lib.SampleHuntMock()
    hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False,
                                 nonadmin_token)

    errors = list(hunt.GetClientsErrors())
    # Make sure there are errors...
    self.assertTrue(errors)
    # but they are not UnauthorizedAccess.
    for e in errors:
      self.assertTrue("UnauthorizedAccess" not in e.backtrace)


class VerifyHuntOutputPluginsCronFlowTest(flow_test_lib.FlowTestsBaseclass,
                                          StandardHuntTestMixin):
  """Tests VerifyHuntOutputPluginsCronFlow."""

  def setUp(self):
    super(VerifyHuntOutputPluginsCronFlowTest, self).setUp()
    self.client_ids = self.SetupClients(1)
    VerifiableDummyHuntOutputPluginVerfier.num_calls = 0

  def GetVerificationsStats(self):
    fields = stats.STATS.GetMetricFields("hunt_output_plugin_verifications")

    result = {}
    for field_value in fields:
      result[field_value[0]] = stats.STATS.GetMetricValue(
          "hunt_output_plugin_verifications", fields=field_value)

    return sum(result.values()), result

  def testDoesNothingWithNonGenericHunts(self):
    implementation.GRRHunt.StartHunt(
        hunt_name=standard.SampleHunt.__name__, token=self.token)

    prev_count, _ = self.GetVerificationsStats()
    flow.GRRFlow.StartFlow(
        flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
        token=self.token)
    count, _ = self.GetVerificationsStats()

    self.assertEqual(count - prev_count, 0)

  def testDoesNothinngWithGenericHuntWithNoPlugins(self):
    self.StartHunt()
    self.AssignTasksToClients()
    self.RunHunt()

    prev_count, _ = self.GetVerificationsStats()
    flow.GRRFlow.StartFlow(
        flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
        token=self.token)
    count, _ = self.GetVerificationsStats()

    self.assertEqual(count - prev_count, 0)

  def testReturnsNAStatusForGenericHuntWithUnverifiableOutputPlugins(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name=DummyHuntOutputPlugin.__name__)
    ])
    self.AssignTasksToClients()
    self.RunHunt()

    _, prev_results = self.GetVerificationsStats()
    flow.GRRFlow.StartFlow(
        flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
        token=self.token)
    _, results = self.GetVerificationsStats()

    self.assertEqual(results["N_A"] - prev_results.get("N_A", 0), 1)

  def testWritesStatusToHuntResultsMetadata(self):
    hunt_urn = self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name=DummyHuntOutputPlugin.__name__)
    ])
    self.AssignTasksToClients()
    self.RunHunt()

    flow.GRRFlow.StartFlow(
        flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
        token=self.token)

    results_metadata = aff4.FACTORY.Open(
        hunt_urn.Add("ResultsMetadata"), token=self.token)
    verification_results_list = results_metadata.Get(
        results_metadata.Schema.OUTPUT_PLUGINS_VERIFICATION_RESULTS)
    self.assertEqual(len(verification_results_list.results), 1)
    self.assertEqual(verification_results_list.results[0].status, "N_A")

  def testRunsCorrespondingVerifierIfThereIsOne(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name=VerifiableDummyHuntOutputPlugin.__name__)
    ])
    self.AssignTasksToClients()
    self.RunHunt()

    _, prev_results = self.GetVerificationsStats()
    flow.GRRFlow.StartFlow(
        flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
        token=self.token)
    _, results = self.GetVerificationsStats()

    self.assertEqual(results["SUCCESS"] - prev_results.get("SUCCESS", 0), 1)

  def testRaisesIfVerifierRaises(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name=DummyHuntOutputPluginWithRaisingVerifier.__name__)
    ])
    self.AssignTasksToClients()
    self.RunHunt()
    prev_count, _ = self.GetVerificationsStats()
    self.assertRaises(
        standard.HuntVerificationError,
        flow.GRRFlow.StartFlow,
        flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
        token=self.token)
    count, _ = self.GetVerificationsStats()
    self.assertEqual(count - prev_count, 0)

  def testUpdatesStatsCounterOnException(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name=DummyHuntOutputPluginWithRaisingVerifier.__name__)
    ])
    self.AssignTasksToClients()
    self.RunHunt()

    prev_count = stats.STATS.GetMetricValue(
        "hunt_output_plugin_verification_errors")
    try:
      flow.GRRFlow.StartFlow(
          flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
          token=self.token)
    except standard.HuntVerificationError:
      pass
    count = stats.STATS.GetMetricValue("hunt_output_plugin_verification_errors")
    self.assertEqual(count - prev_count, 1)

  def testChecksAllHuntsEvenIfOneRaises(self):
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name=DummyHuntOutputPluginWithRaisingVerifier.__name__)
    ])
    self.StartHunt(output_plugins=[
        output_plugin.OutputPluginDescriptor(
            plugin_name=VerifiableDummyHuntOutputPlugin.__name__)
    ])
    self.AssignTasksToClients()
    self.RunHunt()

    _, prev_results = self.GetVerificationsStats()
    self.assertRaises(
        standard.HuntVerificationError,
        flow.GRRFlow.StartFlow,
        flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
        token=self.token)
    _, results = self.GetVerificationsStats()
    self.assertEqual(results["SUCCESS"] - prev_results.get("SUCCESS", 0), 1)

  def testDoesNotCheckHuntsOutsideOfCheckRange(self):
    now = rdfvalue.RDFDatetime.Now()
    with test_lib.FakeTime(now - rdfvalue.Duration("61m"), increment=1e-6):
      self.StartHunt(output_plugins=[
          output_plugin.OutputPluginDescriptor(
              plugin_name=VerifiableDummyHuntOutputPlugin.__name__)
      ])
      self.AssignTasksToClients()
      self.RunHunt()

    prev_count, _ = self.GetVerificationsStats()
    flow.GRRFlow.StartFlow(
        flow_name=standard.VerifyHuntOutputPluginsCronFlow.__name__,
        check_range="60m",
        token=self.token)
    count, _ = self.GetVerificationsStats()
    self.assertEqual(count - prev_count, 0)


def main(argv):
  test_lib.main(argv)


if __name__ == "__main__":
  flags.StartMain(main)
