#!/usr/bin/env python
"""This modules contains tests for hunts API handlers."""

import os
import StringIO
import tarfile
import zipfile


import yaml

from grr.lib import flags

from grr.lib import rdfvalue
from grr.lib import utils

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 test_base as rdf_test_base
from grr.server.grr_response_server import access_control
from grr.server.grr_response_server import aff4
from grr.server.grr_response_server import data_store
from grr.server.grr_response_server import output_plugin
from grr.server.grr_response_server.aff4_objects import aff4_grr
from grr.server.grr_response_server.flows.general import file_finder
from grr.server.grr_response_server.gui import api_test_lib
from grr.server.grr_response_server.gui.api_plugins import hunt as hunt_plugin
from grr.server.grr_response_server.hunts import implementation
from grr.server.grr_response_server.hunts import standard
from grr.server.grr_response_server.output_plugins import test_plugins
from grr.server.grr_response_server.rdfvalues import flow_runner as rdf_flow_runner
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 ApiHuntIdTest(rdf_test_base.RDFValueTestMixin, test_lib.GRRBaseTest):
  """Test for ApiHuntId."""

  rdfvalue_class = hunt_plugin.ApiHuntId

  def GenerateSample(self, number=0):
    return hunt_plugin.ApiHuntId("H:%d" % number)

  def testRaisesWhenInitializedFromInvalidValues(self):
    with self.assertRaises(ValueError):
      hunt_plugin.ApiHuntId("blah")

    with self.assertRaises(ValueError):
      hunt_plugin.ApiHuntId("H:")

    with self.assertRaises(ValueError):
      hunt_plugin.ApiHuntId("1234")

    with self.assertRaises(ValueError):
      hunt_plugin.ApiHuntId("H:1234/foo")

  def testRaisesWhenToURNCalledOnUninitializedValue(self):
    hunt_id = hunt_plugin.ApiHuntId()
    with self.assertRaises(ValueError):
      hunt_id.ToURN()

  def testConvertsToHunttURN(self):
    hunt_id = hunt_plugin.ApiHuntId("H:1234")
    hunt_urn = hunt_id.ToURN()

    self.assertEqual(hunt_urn.Basename(), hunt_id)
    self.assertEqual(hunt_urn, "aff4:/hunts/H:1234")


class ApiCreateHuntHandlerTest(api_test_lib.ApiCallHandlerTest,
                               hunt_test_lib.StandardHuntTestMixin):
  """Test for ApiCreateHuntHandler."""

  def setUp(self):
    super(ApiCreateHuntHandlerTest, self).setUp()
    self.handler = hunt_plugin.ApiCreateHuntHandler()

  def testQueueHuntRunnerArgumentIsNotRespected(self):
    args = hunt_plugin.ApiCreateHuntArgs(
        flow_name=file_finder.FileFinder.__name__)
    args.hunt_runner_args.queue = "BLAH"
    result = self.handler.Handle(args, token=self.token)
    self.assertFalse(result.hunt_runner_args.HasField("queue"))


class ApiListHuntsHandlerTest(api_test_lib.ApiCallHandlerTest,
                              hunt_test_lib.StandardHuntTestMixin):
  """Test for ApiListHuntsHandler."""

  def setUp(self):
    super(ApiListHuntsHandlerTest, self).setUp()
    self.handler = hunt_plugin.ApiListHuntsHandler()

  def testHandlesListOfHuntObjects(self):
    for i in range(10):
      self.CreateHunt(description="hunt_%d" % i)

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(), token=self.token)
    descriptions = set(r.description for r in result.items)

    self.assertEqual(len(descriptions), 10)
    for i in range(10):
      self.assertTrue("hunt_%d" % i in descriptions)

  def testHuntListIsSortedInReversedCreationTimestampOrder(self):
    for i in range(1, 11):
      with test_lib.FakeTime(i * 1000):
        self.CreateHunt(description="hunt_%d" % i)

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(), token=self.token)
    create_times = [r.created.AsMicrosecondsSinceEpoch() for r in result.items]

    self.assertEqual(len(create_times), 10)
    for index, expected_time in enumerate(reversed(range(1, 11))):
      self.assertEqual(create_times[index], expected_time * 1000000000)

  def testHandlesSubrangeOfListOfHuntObjects(self):
    for i in range(1, 11):
      with test_lib.FakeTime(i * 1000):
        self.CreateHunt(description="hunt_%d" % i)

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(offset=2, count=2), token=self.token)
    create_times = [r.created.AsMicrosecondsSinceEpoch() for r in result.items]

    self.assertEqual(len(create_times), 2)
    self.assertEqual(create_times[0], 8 * 1000000000)
    self.assertEqual(create_times[1], 7 * 1000000000)

  def testFiltersHuntsByActivityTime(self):
    for i in range(1, 11):
      with test_lib.FakeTime(i * 60):
        self.CreateHunt(description="hunt_%d" % i)

    with test_lib.FakeTime(10 * 60 + 1):
      result = self.handler.Handle(
          hunt_plugin.ApiListHuntsArgs(active_within="2m"), token=self.token)

    create_times = [r.created for r in result.items]

    self.assertEqual(len(create_times), 2)
    self.assertEqual(create_times[0], 10 * 60 * 1000000)
    self.assertEqual(create_times[1], 9 * 60 * 1000000)

  def testRaisesIfCreatedByFilterUsedWithoutActiveWithinFilter(self):
    self.assertRaises(
        ValueError,
        self.handler.Handle,
        hunt_plugin.ApiListHuntsArgs(created_by="user-bar"),
        token=self.token)

  def testFiltersHuntsByCreator(self):
    for i in range(5):
      self.CreateHunt(
          description="foo_hunt_%d" % i,
          token=access_control.ACLToken(username="user-foo"))

    for i in range(3):
      self.CreateHunt(
          description="bar_hunt_%d" % i,
          token=access_control.ACLToken(username="user-bar"))

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(created_by="user-foo", active_within="1d"),
        token=self.token)
    self.assertEqual(len(result.items), 5)
    for item in result.items:
      self.assertEqual(item.creator, "user-foo")

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(created_by="user-bar", active_within="1d"),
        token=self.token)
    self.assertEqual(len(result.items), 3)
    for item in result.items:
      self.assertEqual(item.creator, "user-bar")

  def testRaisesIfDescriptionContainsFilterUsedWithoutActiveWithinFilter(self):
    self.assertRaises(
        ValueError,
        self.handler.Handle,
        hunt_plugin.ApiListHuntsArgs(description_contains="foo"),
        token=self.token)

  def testFiltersHuntsByDescriptionContainsMatch(self):
    for i in range(5):
      self.CreateHunt(description="foo_hunt_%d" % i)

    for i in range(3):
      self.CreateHunt(description="bar_hunt_%d")

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(
            description_contains="foo", active_within="1d"),
        token=self.token)
    self.assertEqual(len(result.items), 5)
    for item in result.items:
      self.assertTrue("foo" in item.description)

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(
            description_contains="bar", active_within="1d"),
        token=self.token)
    self.assertEqual(len(result.items), 3)
    for item in result.items:
      self.assertTrue("bar" in item.description)

  def testOffsetIsRelativeToFilteredResultsWhenFilterIsPresent(self):
    for i in range(5):
      self.CreateHunt(description="foo_hunt_%d" % i)

    for i in range(3):
      self.CreateHunt(description="bar_hunt_%d" % i)

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(
            description_contains="bar", active_within="1d", offset=1),
        token=self.token)
    self.assertEqual(len(result.items), 2)
    for item in result.items:
      self.assertTrue("bar" in item.description)

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(
            description_contains="bar", active_within="1d", offset=2),
        token=self.token)
    self.assertEqual(len(result.items), 1)
    for item in result.items:
      self.assertTrue("bar" in item.description)

    result = self.handler.Handle(
        hunt_plugin.ApiListHuntsArgs(
            description_contains="bar", active_within="1d", offset=3),
        token=self.token)
    self.assertEqual(len(result.items), 0)


class ApiGetHuntFilesArchiveHandlerTest(api_test_lib.ApiCallHandlerTest,
                                        hunt_test_lib.StandardHuntTestMixin):

  def setUp(self):
    super(ApiGetHuntFilesArchiveHandlerTest, self).setUp()

    self.handler = hunt_plugin.ApiGetHuntFilesArchiveHandler()

    self.hunt = implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flow_runner.FlowRunnerArgs(
            flow_name=file_finder.FileFinder.__name__),
        flow_args=rdf_file_finder.FileFinderArgs(
            paths=[os.path.join(self.base_path, "test.plist")],
            action=rdf_file_finder.FileFinderAction(action_type="DOWNLOAD"),
        ),
        client_rate=0,
        token=self.token)
    self.hunt.Run()

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

  def testGeneratesZipArchive(self):
    result = self.handler.Handle(
        hunt_plugin.ApiGetHuntFilesArchiveArgs(
            hunt_id=self.hunt.urn.Basename(), archive_format="ZIP"),
        token=self.token)

    out_fd = StringIO.StringIO()
    for chunk in result.GenerateContent():
      out_fd.write(chunk)

    zip_fd = zipfile.ZipFile(out_fd, "r")
    manifest = None
    for name in zip_fd.namelist():
      if name.endswith("MANIFEST"):
        manifest = yaml.safe_load(zip_fd.read(name))

    self.assertEqual(manifest["archived_files"], 10)
    self.assertEqual(manifest["failed_files"], 0)
    self.assertEqual(manifest["processed_files"], 10)
    self.assertEqual(manifest["ignored_files"], 0)

  def testGeneratesTarGzArchive(self):
    result = self.handler.Handle(
        hunt_plugin.ApiGetHuntFilesArchiveArgs(
            hunt_id=self.hunt.urn.Basename(), archive_format="TAR_GZ"),
        token=self.token)

    with utils.TempDirectory() as temp_dir:
      tar_path = os.path.join(temp_dir, "archive.tar.gz")
      with open(tar_path, "wb") as fd:
        for chunk in result.GenerateContent():
          fd.write(chunk)

      with tarfile.open(tar_path) as tar_fd:
        tar_fd.extractall(path=temp_dir)

      manifest_file_path = None
      for parent, _, files in os.walk(temp_dir):
        if "MANIFEST" in files:
          manifest_file_path = os.path.join(parent, "MANIFEST")
          break

      self.assertTrue(manifest_file_path)
      with open(manifest_file_path, "rb") as fd:
        manifest = yaml.safe_load(fd.read())

        self.assertEqual(manifest["archived_files"], 10)
        self.assertEqual(manifest["failed_files"], 0)
        self.assertEqual(manifest["processed_files"], 10)
        self.assertEqual(manifest["ignored_files"], 0)


class ApiGetHuntFileHandlerTest(api_test_lib.ApiCallHandlerTest,
                                hunt_test_lib.StandardHuntTestMixin):

  def setUp(self):
    super(ApiGetHuntFileHandlerTest, self).setUp()

    self.handler = hunt_plugin.ApiGetHuntFileHandler()

    self.file_path = os.path.join(self.base_path, "test.plist")
    self.hunt = implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flow_runner.FlowRunnerArgs(
            flow_name=file_finder.FileFinder.__name__),
        flow_args=rdf_file_finder.FileFinderArgs(
            paths=[self.file_path],
            action=rdf_file_finder.FileFinderAction(action_type="DOWNLOAD"),
        ),
        client_rate=0,
        token=self.token)
    self.hunt.Run()

    self.aff4_file_path = "fs/os/%s" % self.file_path

    self.client_id = self.SetupClient(0)
    self.AssignTasksToClients(client_ids=[self.client_id])
    action_mock = action_mocks.FileFinderClientMock()
    hunt_test_lib.TestHuntHelper(
        action_mock, [self.client_id], token=self.token)

  def testRaisesIfOneOfArgumentAttributesIsNone(self):
    model_args = hunt_plugin.ApiGetHuntFileArgs(
        hunt_id=self.hunt.urn.Basename(),
        client_id=self.client_id,
        vfs_path=self.aff4_file_path,
        timestamp=rdfvalue.RDFDatetime.Now())

    with self.assertRaises(ValueError):
      args = model_args.Copy()
      args.hunt_id = None
      self.handler.Handle(args)

    with self.assertRaises(ValueError):
      args = model_args.Copy()
      args.client_id = None
      self.handler.Handle(args)

    with self.assertRaises(ValueError):
      args = model_args.Copy()
      args.vfs_path = None
      self.handler.Handle(args)

    with self.assertRaises(ValueError):
      args = model_args.Copy()
      args.timestamp = None
      self.handler.Handle(args)

  def testRaisesIfVfsRootIsNotWhitelisted(self):
    args = hunt_plugin.ApiGetHuntFileArgs(
        hunt_id=self.hunt.urn.Basename(),
        client_id=self.client_id,
        vfs_path="flows/W:123456",
        timestamp=rdfvalue.RDFDatetime().Now())

    with self.assertRaises(ValueError):
      self.handler.Handle(args)

  def testRaisesIfResultIsBeforeTimestamp(self):
    results = implementation.GRRHunt.ResultCollectionForHID(
        self.hunt.urn, token=self.token)

    args = hunt_plugin.ApiGetHuntFileArgs(
        hunt_id=self.hunt.urn.Basename(),
        client_id=self.client_id,
        vfs_path=self.aff4_file_path,
        timestamp=results[0].age + rdfvalue.Duration("1s"))
    with self.assertRaises(hunt_plugin.HuntFileNotFoundError):
      self.handler.Handle(args, token=self.token)

  def _FillInStubResults(self):
    results = implementation.GRRHunt.ResultCollectionForHID(
        self.hunt.urn, token=self.token)
    result = results[0]

    with data_store.DB.GetMutationPool() as pool:
      for i in range(self.handler.MAX_RECORDS_TO_CHECK):
        wrong_result = rdf_flows.GrrMessage(
            payload=rdfvalue.RDFString("foo/bar"),
            age=(result.age - (self.handler.MAX_RECORDS_TO_CHECK - i + 1) *
                 rdfvalue.Duration("1s")),
            source=self.client_id)
        results.Add(
            wrong_result, timestamp=wrong_result.age, mutation_pool=pool)

    return result

  def testRaisesIfResultIsAfterMaxRecordsAfterTimestamp(self):
    original_result = self._FillInStubResults()

    args = hunt_plugin.ApiGetHuntFileArgs(
        hunt_id=self.hunt.urn.Basename(),
        client_id=self.client_id,
        vfs_path=self.aff4_file_path,
        timestamp=original_result.age -
        (self.handler.MAX_RECORDS_TO_CHECK + 1) * rdfvalue.Duration("1s"))

    with self.assertRaises(hunt_plugin.HuntFileNotFoundError):
      self.handler.Handle(args, token=self.token)

  def testReturnsResultIfWithinMaxRecordsAfterTimestamp(self):
    original_result = self._FillInStubResults()

    args = hunt_plugin.ApiGetHuntFileArgs(
        hunt_id=self.hunt.urn.Basename(),
        client_id=self.client_id,
        vfs_path=self.aff4_file_path,
        timestamp=original_result.age -
        self.handler.MAX_RECORDS_TO_CHECK * rdfvalue.Duration("1s"))

    self.handler.Handle(args, token=self.token)

  def testRaisesIfResultFileIsNotStream(self):
    original_results = implementation.GRRHunt.ResultCollectionForHID(
        self.hunt.urn, token=self.token)
    original_result = original_results[0]

    with aff4.FACTORY.Create(
        original_result.payload.stat_entry.AFF4Path(self.client_id),
        aff4_type=aff4.AFF4Volume,
        token=self.token):
      pass

    args = hunt_plugin.ApiGetHuntFileArgs(
        hunt_id=self.hunt.urn.Basename(),
        client_id=self.client_id,
        vfs_path=self.aff4_file_path,
        timestamp=original_result.age)

    with self.assertRaises(hunt_plugin.HuntFileNotFoundError):
      self.handler.Handle(args, token=self.token)

  def testRaisesIfResultIsEmptyStream(self):
    original_results = implementation.GRRHunt.ResultCollectionForHID(
        self.hunt.urn, token=self.token)
    original_result = original_results[0]

    urn = original_result.payload.stat_entry.AFF4Path(self.client_id)
    aff4.FACTORY.Delete(urn, token=self.token)
    with aff4.FACTORY.Create(urn, aff4_type=aff4_grr.VFSFile, token=self.token):
      pass

    args = hunt_plugin.ApiGetHuntFileArgs(
        hunt_id=self.hunt.urn.Basename(),
        client_id=self.client_id,
        vfs_path=self.aff4_file_path,
        timestamp=original_result.age)

    with self.assertRaises(hunt_plugin.HuntFileNotFoundError):
      self.handler.Handle(args, token=self.token)

  def testReturnsBinaryStreamIfResultFound(self):
    results = implementation.GRRHunt.ResultCollectionForHID(
        self.hunt.urn, token=self.token)

    args = hunt_plugin.ApiGetHuntFileArgs(
        hunt_id=self.hunt.urn.Basename(),
        client_id=self.client_id,
        vfs_path=self.aff4_file_path,
        timestamp=results[0].age)

    result = self.handler.Handle(args, token=self.token)
    self.assertTrue(hasattr(result, "GenerateContent"))
    self.assertEqual(result.content_length,
                     results[0].payload.stat_entry.st_size)


class ApiListHuntOutputPluginLogsHandlerTest(
    api_test_lib.ApiCallHandlerTest, hunt_test_lib.StandardHuntTestMixin):
  """Test for ApiListHuntOutputPluginLogsHandler."""

  def setUp(self):
    super(ApiListHuntOutputPluginLogsHandlerTest, self).setUp()

    self.client_ids = self.SetupClients(5)
    self.handler = hunt_plugin.ApiListHuntOutputPluginLogsHandler()
    self.output_plugins = [
        output_plugin.OutputPluginDescriptor(
            plugin_name=test_plugins.DummyHuntTestOutputPlugin.__name__,
            plugin_args=test_plugins.DummyHuntTestOutputPlugin.args_type(
                filename_regex="foo")),
        output_plugin.OutputPluginDescriptor(
            plugin_name=test_plugins.DummyHuntTestOutputPlugin.__name__,
            plugin_args=test_plugins.DummyHuntTestOutputPlugin.args_type(
                filename_regex="bar"))
    ]

  def RunHuntWithOutputPlugins(self, output_plugins):
    hunt_urn = self.StartHunt(
        description="the hunt", output_plugins=output_plugins)

    for client_id in self.client_ids:
      self.AssignTasksToClients(client_ids=[client_id])
      self.RunHunt(failrate=-1)
      self.ProcessHuntOutputPlugins()

    return hunt_urn

  def testReturnsLogsWhenJustOnePlugin(self):
    hunt_urn = self.RunHuntWithOutputPlugins([self.output_plugins[0]])
    result = self.handler.Handle(
        hunt_plugin.ApiListHuntOutputPluginLogsArgs(
            hunt_id=hunt_urn.Basename(),
            plugin_id=test_plugins.DummyHuntTestOutputPlugin.__name__ + "_0"),
        token=self.token)

    self.assertEqual(result.total_count, 5)
    self.assertEqual(len(result.items), 5)
    for item in result.items:
      self.assertEqual("foo", item.plugin_descriptor.plugin_args.filename_regex)

  def testReturnsLogsWhenMultiplePlugins(self):
    hunt_urn = self.RunHuntWithOutputPlugins(self.output_plugins)
    result = self.handler.Handle(
        hunt_plugin.ApiListHuntOutputPluginLogsArgs(
            hunt_id=hunt_urn.Basename(),
            plugin_id=test_plugins.DummyHuntTestOutputPlugin.__name__ + "_1"),
        token=self.token)

    self.assertEqual(result.total_count, 5)
    self.assertEqual(len(result.items), 5)
    for item in result.items:
      self.assertEqual("bar", item.plugin_descriptor.plugin_args.filename_regex)

  def testSlicesLogsWhenJustOnePlugin(self):
    hunt_urn = self.RunHuntWithOutputPlugins([self.output_plugins[0]])
    result = self.handler.Handle(
        hunt_plugin.ApiListHuntOutputPluginLogsArgs(
            hunt_id=hunt_urn.Basename(),
            offset=2,
            count=2,
            plugin_id=test_plugins.DummyHuntTestOutputPlugin.__name__ + "_0"),
        token=self.token)

    self.assertEqual(result.total_count, 5)
    self.assertEqual(len(result.items), 2)

  def testSlicesLogsWhenMultiplePlugins(self):
    hunt_urn = self.RunHuntWithOutputPlugins(self.output_plugins)
    result = self.handler.Handle(
        hunt_plugin.ApiListHuntOutputPluginLogsArgs(
            hunt_id=hunt_urn.Basename(),
            offset=2,
            count=2,
            plugin_id=test_plugins.DummyHuntTestOutputPlugin.__name__ + "_1"),
        token=self.token)

    self.assertEqual(result.total_count, 5)
    self.assertEqual(len(result.items), 2)


class ApiModifyHuntHandlerTest(api_test_lib.ApiCallHandlerTest,
                               hunt_test_lib.StandardHuntTestMixin):
  """Test for ApiModifyHuntHandler."""

  def setUp(self):
    super(ApiModifyHuntHandlerTest, self).setUp()

    self.handler = hunt_plugin.ApiModifyHuntHandler()

    self.hunt = self.CreateHunt(description="the hunt")
    self.hunt_urn = self.hunt.urn

    self.args = hunt_plugin.ApiModifyHuntArgs(hunt_id=self.hunt.urn.Basename())

  def testDoesNothingIfArgsHaveNoChanges(self):
    before = hunt_plugin.ApiHunt().InitFromAff4Object(
        aff4.FACTORY.Open(self.hunt.urn, token=self.token))

    self.handler.Handle(self.args, token=self.token)

    after = hunt_plugin.ApiHunt().InitFromAff4Object(
        aff4.FACTORY.Open(self.hunt.urn, token=self.token))

    self.assertEqual(before, after)

  def testRaisesIfStateIsSetToNotStartedOrStopped(self):
    with self.assertRaises(hunt_plugin.InvalidHuntStateError):
      self.args.state = "COMPLETED"
      self.handler.Handle(self.args, token=self.token)

  def testRaisesWhenStartingHuntInTheWrongState(self):
    self.hunt.Run()
    self.hunt.Stop()

    with self.assertRaises(hunt_plugin.HuntNotStartableError):
      self.args.state = "STARTED"
      self.handler.Handle(self.args, token=self.token)

  def testRaisesWhenStoppingHuntInTheWrongState(self):
    self.hunt.Run()
    self.hunt.Stop()

    with self.assertRaises(hunt_plugin.HuntNotStoppableError):
      self.args.state = "STOPPED"
      self.handler.Handle(self.args, token=self.token)

  def testStartsHuntCorrectly(self):
    self.args.state = "STARTED"
    self.handler.Handle(self.args, token=self.token)

    hunt = aff4.FACTORY.Open(self.hunt_urn, token=self.token)
    self.assertEqual(hunt.Get(hunt.Schema.STATE), "STARTED")

  def testStopsHuntCorrectly(self):
    self.args.state = "STOPPED"
    self.handler.Handle(self.args, token=self.token)

    hunt = aff4.FACTORY.Open(self.hunt_urn, token=self.token)
    self.assertEqual(hunt.Get(hunt.Schema.STATE), "STOPPED")

  def testRaisesWhenModifyingHuntInNonPausedState(self):
    self.hunt.Run()

    with self.assertRaises(hunt_plugin.HuntNotModifiableError):
      self.args.client_rate = 100
      self.handler.Handle(self.args, token=self.token)

  def testModifiesHuntCorrectly(self):
    self.args.client_rate = 100
    self.args.client_limit = 42
    self.args.expires = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(42)

    self.handler.Handle(self.args, token=self.token)

    after = hunt_plugin.ApiHunt().InitFromAff4Object(
        aff4.FACTORY.Open(self.hunt.urn, token=self.token))
    self.assertEqual(after.client_rate, 100)
    self.assertEqual(after.client_limit, 42)
    self.assertEqual(after.expires,
                     rdfvalue.RDFDatetime.FromSecondsSinceEpoch(42))

  def testDoesNotModifyHuntIfStateChangeFails(self):
    with self.assertRaises(hunt_plugin.InvalidHuntStateError):
      self.args.client_limit = 42
      self.args.state = "COMPLETED"
      self.handler.Handle(self.args, token=self.token)

    after = hunt_plugin.ApiHunt().InitFromAff4Object(
        aff4.FACTORY.Open(self.hunt.urn, token=self.token))
    self.assertNotEqual(after.client_limit, 42)


class ApiDeleteHuntHandlerTest(api_test_lib.ApiCallHandlerTest,
                               hunt_test_lib.StandardHuntTestMixin):
  """Test for ApiDeleteHuntHandler."""

  def setUp(self):
    super(ApiDeleteHuntHandlerTest, self).setUp()

    self.handler = hunt_plugin.ApiDeleteHuntHandler()

    self.hunt = self.CreateHunt(description="the hunt")
    self.hunt_urn = self.hunt.urn

    self.args = hunt_plugin.ApiDeleteHuntArgs(hunt_id=self.hunt.urn.Basename())

  def testRaisesIfHuntNotFound(self):
    with self.assertRaises(hunt_plugin.HuntNotFoundError):
      self.handler.Handle(
          hunt_plugin.ApiDeleteHuntArgs(hunt_id="H:123456"), token=self.token)

  def testRaisesIfHuntIsRunning(self):
    self.hunt.Run()

    with self.assertRaises(hunt_plugin.HuntNotDeletableError):
      self.handler.Handle(self.args, token=self.token)

  def testDeletesHunt(self):
    self.handler.Handle(self.args, token=self.token)

    with self.assertRaises(aff4.InstantiationError):
      aff4.FACTORY.Open(
          self.hunt_urn, aff4_type=implementation.GRRHunt, token=self.token)


class ApiGetExportedHuntResultsHandlerTest(test_lib.GRRBaseTest,
                                           hunt_test_lib.StandardHuntTestMixin):

  def setUp(self):
    super(ApiGetExportedHuntResultsHandlerTest, self).setUp()

    self.handler = hunt_plugin.ApiGetExportedHuntResultsHandler()

    self.hunt = implementation.GRRHunt.StartHunt(
        hunt_name=standard.GenericHunt.__name__,
        flow_runner_args=rdf_flow_runner.FlowRunnerArgs(
            flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__),
        client_rate=0,
        token=self.token)
    self.hunt.Run()

    self.client_ids = self.SetupClients(5)
    # Ensure that clients are processed sequentially - this way the test won't
    # depend on the order of results in the collection (which is normally
    # random).
    for cid in self.client_ids:
      self.AssignTasksToClients(client_ids=[cid])
      client_mock = hunt_test_lib.SampleHuntMock()
      hunt_test_lib.TestHuntHelper(client_mock, [cid], token=self.token)

  def testWorksCorrectlyWithTestOutputPluginOnFlowWithSingleResult(self):
    result = self.handler.Handle(
        hunt_plugin.ApiGetExportedHuntResultsArgs(
            hunt_id=self.hunt.urn.Basename(),
            plugin_name=test_plugins.TestInstantOutputPlugin.plugin_name),
        token=self.token)

    chunks = list(result.GenerateContent())

    self.assertListEqual(
        chunks,
        ["Start: %s" % utils.SmartStr(self.hunt.urn),
         "Values of type: RDFString",
         "First pass: oh (source=%s)" % self.client_ids[0],
         "First pass: oh (source=%s)" % self.client_ids[1],
         "First pass: oh (source=%s)" % self.client_ids[2],
         "First pass: oh (source=%s)" % self.client_ids[3],
         "First pass: oh (source=%s)" % self.client_ids[4],
         "Second pass: oh (source=%s)" % self.client_ids[0],
         "Second pass: oh (source=%s)" % self.client_ids[1],
         "Second pass: oh (source=%s)" % self.client_ids[2],
         "Second pass: oh (source=%s)" % self.client_ids[3],
         "Second pass: oh (source=%s)" % self.client_ids[4],
         "Finish: %s" % utils.SmartStr(self.hunt.urn)])  # pyformat: disable


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


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