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




from grr.gui import api_regression_test_lib
from grr.gui.api_plugins import client as client_plugin

from grr.lib import aff4
from grr.lib import flags
from grr.lib import flow
from grr.lib import queue_manager
from grr.lib import rdfvalue
from grr.lib import test_lib

from grr.lib.aff4_objects import aff4_grr
from grr.lib.aff4_objects import stats as aff4_stats
from grr.lib.flows.general import processes
from grr.lib.hunts import standard_test
from grr.lib.rdfvalues import client as rdf_client


class ApiSearchClientsHandlerRegressionTest(
    api_regression_test_lib.ApiRegressionTest):

  api_method = "SearchClients"
  handler = client_plugin.ApiSearchClientsHandler

  def Run(self):
    # Fix the time to avoid regressions.
    with test_lib.FakeTime(42):
      client_ids = self.SetupClients(1)

      # Delete the certificate as it's being regenerated every time the
      # client is created.
      with aff4.FACTORY.Open(
          client_ids[0], mode="rw", token=self.token) as grr_client:
        grr_client.DeleteAttribute(grr_client.Schema.CERT)

      self.Check(
          "SearchClients",
          args=client_plugin.ApiSearchClientsArgs(
              query=client_ids[0].Basename()))


class ApiGetClientHandlerRegressionTest(
    api_regression_test_lib.ApiRegressionTest):

  api_method = "GetClient"
  handler = client_plugin.ApiGetClientHandler

  def Run(self):
    # Fix the time to avoid regressions.
    with test_lib.FakeTime(42):
      client_ids = self.SetupClients(1)

      with aff4.FACTORY.Open(
          client_ids[0], mode="rw", token=self.token) as grr_client:
        grr_client.Set(grr_client.Schema.MEMORY_SIZE(4294967296))
        # Delete the certificate as it's being regenerated every time the
        # client is created.
        grr_client.DeleteAttribute(grr_client.Schema.CERT)

    self.Check(
        "GetClient",
        args=client_plugin.ApiGetClientArgs(client_id=client_ids[0].Basename()))


class ApiGetLastClientIPAddressHandlerRegressionTest(
    api_regression_test_lib.ApiRegressionTest):

  api_method = "GetLastClientIPAddress"
  handler = client_plugin.ApiGetLastClientIPAddressHandler

  def Run(self):
    # Fix the time to avoid regressions.
    with test_lib.FakeTime(42):
      client_id = self.SetupClients(1)[0]

      with aff4.FACTORY.Open(
          client_id, mode="rw", token=self.token) as grr_client:
        grr_client.Set(grr_client.Schema.CLIENT_IP("192.168.100.42"))

    self.Check(
        "GetLastClientIPAddress",
        args=client_plugin.ApiGetLastClientIPAddressArgs(
            client_id=client_id.Basename()))


class ApiListClientsLabelsHandlerRegressionTest(
    api_regression_test_lib.ApiRegressionTest):

  api_method = "ListClientsLabels"
  handler = client_plugin.ApiListClientsLabelsHandler

  def Run(self):
    # Fix the time to avoid regressions.
    with test_lib.FakeTime(42):
      client_ids = self.SetupClients(2)

      with aff4.FACTORY.Open(
          client_ids[0], mode="rw", token=self.token) as grr_client:
        grr_client.AddLabels("foo")

      with aff4.FACTORY.Open(
          client_ids[1], mode="rw", token=self.token) as grr_client:
        grr_client.AddLabels("bar")

    self.Check("ListClientsLabels")


class ApiListKbFieldsHandlerTest(api_regression_test_lib.ApiRegressionTest):

  api_method = "ListKbFields"
  handler = client_plugin.ApiListKbFieldsHandler

  def Run(self):
    self.Check("ListKbFields")


class ApiListClientCrashesHandlerRegressionTest(
    api_regression_test_lib.ApiRegressionTest,
    standard_test.StandardHuntTestMixin):

  api_method = "ListClientCrashes"
  handler = client_plugin.ApiListClientCrashesHandler

  def Run(self):
    client_ids = self.SetupClients(1)
    client_id = client_ids[0]
    client_mock = test_lib.CrashClientMock(client_id, self.token)

    with test_lib.FakeTime(42):
      with self.CreateHunt(description="the hunt") as hunt_obj:
        hunt_obj.Run()

    with test_lib.FakeTime(45):
      self.AssignTasksToClients(client_ids)
      test_lib.TestHuntHelperWithMultipleMocks({
          client_id: client_mock
      }, False, self.token)

    crashes = aff4_grr.VFSGRRClient.CrashCollectionForCID(
        client_id, token=self.token)
    crash = list(crashes)[0]