#!/usr/bin/env python
"""Unittest for grr frontend server."""




# pylint: disable=unused-import,g-bad-import-order
from grr.lib import server_plugins
# pylint: enable=unused-import,g-bad-import-order

from grr.lib import communicator
from grr.lib import config_lib
from grr.lib import data_store
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 import utils


class SendingTestFlow(flow.GRRFlow):
  """Tests that sent messages are correctly collected."""

  @flow.StateHandler(next_state="Incoming")
  def Start(self):
    for i in range(10):
      self.CallClient("Test",
                      rdfvalue.DataBlob(string="test%s" % i),
                      data=str(i),
                      next_state="Incoming")


class GRRFEServerTest(test_lib.FlowTestsBaseclass):
  """Tests the GRRFEServer."""
  string = "Test String"

  def setUp(self):
    """Setup the server."""
    super(GRRFEServerTest, self).setUp()

    # Whitelist test flow.
    config_lib.CONFIG.Set("Frontend.well_known_flows", [utils.SmartStr(
        test_lib.WellKnownSessionTest.well_known_session_id)])

    # For tests, small pools are ok.
    config_lib.CONFIG.Set("Threadpool.size", 10)

    prefix = "pool-%s" % self._testMethodName
    self.server = flow.FrontEndServer(
        certificate=config_lib.CONFIG["Frontend.certificate"],
        private_key=config_lib.CONFIG["PrivateKeys.server_key"],
        threadpool_prefix=prefix)

  def CheckMessages(self, left, right):
    """Compares two lists of messages for equality.

    Args:
      left: A list of GrrMessage
      right: A list of (task, GrrMessage)

    Returns:
      True if they are the same.
    """
    if len(right) != len(left):
      return False

    for i in range(len(right)):
      if left[i] != right[i][1]:
        return False

    return True

  def testReceiveMessages(self):
    """Test Receiving messages with no status."""
    flow_obj = self.FlowSetup("FlowOrderTest")

    session_id = flow_obj.session_id
    messages = [rdfvalue.GrrMessage(request_id=1,
                                    response_id=i,
                                    session_id=session_id,
                                    args=str(i))
                for i in range(1, 10)]

    self.server.ReceiveMessages(self.client_id, messages)

    # Make sure the task is still on the client queue
    manager = queue_manager.QueueManager(token=self.token)
    tasks_on_client_queue = manager.Query(self.client_id, 100)
    self.assertEqual(len(tasks_on_client_queue), 1)

    # Check that messages were stored correctly
    for message in messages:
      stored_message, _ = data_store.DB.Resolve(
          session_id.Add("state/request:00000001"),
          manager.FLOW_RESPONSE_TEMPLATE % (1, message.response_id),
          token=self.token)

      stored_message = rdfvalue.GrrMessage(stored_message)
      self.assertProtoEqual(stored_message, message)

    return messages

  def testReceiveMessagesWithStatus(self):
    """Receiving a sequence of messages with a status."""
    flow_obj = self.FlowSetup("FlowOrderTest")

    session_id = flow_obj.session_id
    messages = [rdfvalue.GrrMessage(request_id=1,
                                    response_id=i,
                                    session_id=session_id,
                                    args=str(i),
                                    task_id=15)
                for i in range(1, 10)]

    # Now add the status message
    status = rdfvalue.GrrStatus(status=rdfvalue.GrrStatus.ReturnedStatus.OK)
    messages.append(rdfvalue.GrrMessage(
        request_id=1, response_id=len(messages)+1, task_id=15,
        session_id=messages[0].session_id, payload=status,
        type=rdfvalue.GrrMessage.Type.STATUS))

    self.server.ReceiveMessages(self.client_id, messages)

    # Make sure the task is still on the client queue
    manager = queue_manager.QueueManager(token=self.token)
    tasks_on_client_queue = manager.Query(self.client_id, 100)
    self.assertEqual(len(tasks_on_client_queue), 1)

    # Check that messages were stored correctly
    for message in messages:
      stored_message, _ = data_store.DB.Resolve(
          session_id.Add("state/request:00000001"),
          manager.FLOW_RESPONSE_TEMPLATE % (1, message.response_id),
          token=self.token)

      stored_message = rdfvalue.GrrMessage(stored_message)
      self.assertProtoEqual(stored_message, message)

  def testWellKnownFlows(self):
    """Make sure that well known flows can run on the front end."""
    test_lib.WellKnownSessionTest.messages = []
    session_id = test_lib.WellKnownSessionTest.well_known_session_id

    messages = [rdfvalue.GrrMessage(request_id=0,
                                    response_id=0,
                                    session_id=session_id,
                                    args=str(i))
                for i in range(1, 10)]

    self.server.ReceiveMessages(self.client_id, messages)

    # Wait for async actions to complete
    self.server.thread_pool.Join()

    test_lib.WellKnownSessionTest.messages.sort()

    # Well known flows are now directly processed on the front end
    self.assertEqual(test_lib.WellKnownSessionTest.messages,
                     list(range(1, 10)))

    # There should be nothing in the client_queue
    self.assertEqual([], data_store.DB.ResolveRegex(self.client_id, "task:.*",
                                                    token=self.token))

  def testWellKnownFlowsRemote(self):
    """Make sure that flows that do not exist on the front end get scheduled."""
    test_lib.WellKnownSessionTest.messages = []
    session_id = test_lib.WellKnownSessionTest.well_known_session_id

    messages = [rdfvalue.GrrMessage(request_id=0,
                                    response_id=0,
                                    session_id=session_id,
                                    args=str(i))
                for i in range(1, 10)]

    # Delete the local well known flow cache is empty.
    self.server.well_known_flows = {}
    self.server.ReceiveMessages(self.client_id, messages)

    # Wait for async actions to complete
    self.server.thread_pool.Join()

    # None get processed now
    self.assertEqual(test_lib.WellKnownSessionTest.messages, [])

    # There should be nothing in the client_queue
    self.assertEqual([], data_store.DB.ResolveRegex(self.client_id, "task:.*",
                                                    token=self.token))

    # The well known flow messages should be waiting in the flow state now:
    queued_messages = []
    for predicate, _, _ in data_store.DB.ResolveRegex(
        session_id.Add("state/request:00000000"), "flow:.*", token=self.token):
      queued_messages.append(predicate)

    self.assertEqual(len(queued_messages), 9)

  def testDrainUpdateSessionRequestStates(self):
    """Draining the flow requests and preparing messages."""
    # This flow sends 10 messages on Start()
    flow_obj = self.FlowSetup("SendingTestFlow")
    session_id = flow_obj.session_id

    # There should be 10 messages in the client's task queue
    manager = queue_manager.QueueManager(token=self.token)
    tasks = manager.Query(self.client_id, 100)
    self.assertEqual(len(tasks), 10)

    # Check that the response state objects have the correct ts_id set
    # in the client_queue:
    for task in tasks:
      request_id = task.request_id

      # Retrieve the request state for this request_id
      request_state, _ = data_store.DB.Resolve(
          session_id.Add("state"),
          manager.FLOW_REQUEST_TEMPLATE % request_id,
          token=self.token)

      request_state = rdfvalue.RequestState(request_state)

      # Check that task_id for the client message is correctly set in
      # request_state.
      self.assertEqual(request_state.request.task_id, task.task_id)

    # Now ask the server to drain the outbound messages into the
    # message list.
    response = rdfvalue.MessageList()

    self.server.DrainTaskSchedulerQueueForClient(
        self.client_id, 5, response)

    # Check that we received only as many messages as we asked for
    self.assertEqual(len(response.job), 5)

    for i in range(4):
      self.assertEqual(response.job[i].session_id, session_id)
      self.assertEqual(response.job[i].name, "Test")

  def testUpdateAndCheckIfShouldThrottle(self):
    self.server.SetThrottleBundlesRatio(1.0)

    # Let's assume that requests are flowing in every 10 seconds
    self.server.UpdateAndCheckIfShouldThrottle(0)
    self.server.UpdateAndCheckIfShouldThrottle(10)
    self.server.UpdateAndCheckIfShouldThrottle(20)
    self.server.UpdateAndCheckIfShouldThrottle(30)
    self.server.UpdateAndCheckIfShouldThrottle(40)
    self.server.UpdateAndCheckIfShouldThrottle(50)
    self.server.UpdateAndCheckIfShouldThrottle(60)

    self.server.SetThrottleBundlesRatio(0.3)

    # Now: average interval between requests is 10 seconds
    # According to throttling logic, requests will only be allowed if
    # the interval between them is 10 / 0.3 = 33.3 seconds
    result = self.server.UpdateAndCheckIfShouldThrottle(70)
    self.assertEquals(result, True)

    result = self.server.UpdateAndCheckIfShouldThrottle(80)
    self.assertEquals(result, True)

    result = self.server.UpdateAndCheckIfShouldThrottle(90)
    self.assertEquals(result, True)

    result = self.server.UpdateAndCheckIfShouldThrottle(100)
    self.assertEquals(result, False)

    result = self.server.UpdateAndCheckIfShouldThrottle(110)
    self.assertEquals(result, True)

    result = self.server.UpdateAndCheckIfShouldThrottle(120)
    self.assertEquals(result, True)

    result = self.server.UpdateAndCheckIfShouldThrottle(130)
    self.assertEquals(result, True)

    result = self.server.UpdateAndCheckIfShouldThrottle(140)
    self.assertEquals(result, False)

    # Now we throttle everything
    self.server.SetThrottleBundlesRatio(0)

    result = self.server.UpdateAndCheckIfShouldThrottle(141)
    self.assertEquals(result, True)

    result = self.server.UpdateAndCheckIfShouldThrottle(142)
    self.assertEquals(result, True)