#!/usr/bin/env python
"""Unittest for GRR<->Fleetspeak server side glue code."""

import itertools
import mock

from fleetspeak.src.common.proto.fleetspeak import common_pb2 as fs_common_pb2
from fleetspeak.src.server.grpcservice.client import client as fs_client

from grr.lib import communicator
from grr.lib import flags
from grr.lib import rdfvalue
from grr.lib.rdfvalues import client as rdf_client
from grr.lib.rdfvalues import flows as rdf_flows
from grr_response_proto import jobs_pb2
from grr.server import aff4
from grr.server import data_store
from grr.server import fleetspeak_connector
from grr.server import fleetspeak_utils
from grr.server import flow
from grr.server import front_end_test
from grr.server import queue_manager
from grr.server.flows.general import processes as flow_processes
from grr.test_lib import action_mocks
from grr.test_lib import flow_test_lib
from grr.test_lib import test_lib
from grr.tools import fleetspeak_frontend as fs_frontend_tool


class _FakeGRPCServiceClient(fs_client.ServiceClient):

  def __init__(self, service_name, send_callback=lambda _: None):
    super(_FakeGRPCServiceClient, self).__init__(service_name)
    self._process = None
    self._send_callback = send_callback

  def Listen(self, process):
    self._process = process

  def Send(self, message):
    self._send_callback(message)

  def MockIncomingFSMessage(self, message):
    self._process(message)

  def MockIncomingFSMessages(self, messages):
    for message in messages:
      self.MockIncomingFSMessage(message)


def SetAFF4FSEnabledFlag(grr_id, token):
  with aff4.FACTORY.Create(
      grr_id, aff4.AFF4Object.classes["VFSGRRClient"], mode="w",
      token=token) as client:
    client.Set(client.Schema.FLEETSPEAK_ENABLED, rdfvalue.RDFBool(True))


class FleetspeakGRRFEServerTest(front_end_test.GRRFEServerTestBase):
  """Tests the Fleetspeak based GRRFEServer."""

  def testReceiveMessagesFleetspeak(self):
    service_name = "GRR"
    fake_service_client = _FakeGRPCServiceClient(service_name)

    fleetspeak_connector.Reset()
    fleetspeak_connector.Init(service_client=fake_service_client)

    fsd = fs_frontend_tool.GRRFSServer()

    grr_client_nr = 0xab
    grr_client_id_urn = self.SetupClient(grr_client_nr)

    flow_obj = self.FlowSetup(flow_test_lib.FlowOrderTest.__name__,
                              grr_client_id_urn)

    num_msgs = 9

    session_id = flow_obj.session_id
    messages = [
        rdf_flows.GrrMessage(
            request_id=1,
            response_id=i,
            session_id=session_id,
            payload=rdfvalue.RDFInteger(i)) for i in xrange(1, num_msgs + 1)
    ]

    fs_client_id = "\x10\x00\x00\x00\x00\x00\x00\xab"
    # fs_client_id should be equivalent to grr_client_id_urn
    self.assertEqual(fs_client_id,
                     fleetspeak_utils.GRRIDToFleetspeakID(
                         grr_client_id_urn.Basename()))

    fs_messages = [
        fs_common_pb2.Message(
            message_type="GrrMessage",
            source=fs_common_pb2.Address(
                client_id=fs_client_id, service_name=service_name))
        for _ in xrange(num_msgs)
    ]
    for fs_message, message in itertools.izip(fs_messages, messages):
      fs_message.data.Pack(message.AsPrimitiveProto())

    for msg in fs_messages:
      fsd.Process(msg, None)

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

    want_messages = [message.Copy() for message in messages]
    for want_message in want_messages:
      # This is filled in by the frontend as soon as it gets the message.
      want_message.auth_state = (
          rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED)
      want_message.source = grr_client_id_urn

    stored_messages = data_store.DB.ReadResponsesForRequestId(session_id, 1)

    self.assertEqual(len(stored_messages), len(want_messages))

    stored_messages.sort(key=lambda m: m.response_id)
    # Check that messages were stored correctly
    for stored_message, want_message in itertools.izip(stored_messages,
                                                       want_messages):
      stored_message.timestamp = None
      self.assertRDFValuesEqual(stored_message, want_message)

  def testReceiveMessageListFleetspeak(self):
    service_name = "GRR"
    fake_service_client = _FakeGRPCServiceClient(service_name)

    fleetspeak_connector.Reset()
    fleetspeak_connector.Init(service_client=fake_service_client)

    fsd = fs_frontend_tool.GRRFSServer()

    grr_client_nr = 0xab
    grr_client_id_urn = self.SetupClient(grr_client_nr)

    flow_obj = self.FlowSetup(flow_test_lib.FlowOrderTest.__name__,
                              grr_client_id_urn)

    num_msgs = 9

    session_id = flow_obj.session_id
    messages = [
        rdf_flows.GrrMessage(
            request_id=1,
            response_id=i,
            session_id=session_id,
            payload=rdfvalue.RDFInteger(i)) for i in xrange(1, num_msgs + 1)
    ]

    fs_client_id = "\x10\x00\x00\x00\x00\x00\x00\xab"
    # fs_client_id should be equivalent to grr_client_id_urn
    self.assertEqual(fs_client_id,
                     fleetspeak_utils.GRRIDToFleetspeakID(
                         grr_client_id_urn.Basename()))

    message_list = rdf_flows.PackedMessageList()
    communicator.Communicator.EncodeMessageList(
        rdf_flows.MessageList(job=messages), message_list)

    fs_message = fs_common_pb2.Message(
        message_type="MessageList",
        source=fs_common_pb2.Address(
            client_id=fs_client_id, service_name=service_name))
    fs_message.data.Pack(message_list.AsPrimitiveProto())
    fsd.Process(fs_message, None)

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

    want_messages = [message.Copy() for message in messages]
    for want_message in want_messages:
      # This is filled in by the frontend as soon as it gets the message.
      want_message.auth_state = (
          rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED)
      want_message.source = grr_client_id_urn

    stored_messages = data_store.DB.ReadResponsesForRequestId(session_id, 1)

    self.assertEqual(len(stored_messages), len(want_messages))

    stored_messages.sort(key=lambda m: m.response_id)
    # Check that messages were stored correctly
    for stored_message, want_message in itertools.izip(stored_messages,
                                                       want_messages):
      stored_message.timestamp = None
      self.assertRDFValuesEqual(stored_message, want_message)


class ListProcessesMock(action_mocks.FileFinderClientMock):
  """Client with real file actions and mocked-out ListProcesses."""

  def __init__(self, processes_list):
    super(ListProcessesMock, self).__init__()
    self.processes_list = processes_list

    self.mock_task_queue = []

  def ListProcesses(self, _):
    return self.processes_list


class ListProcessesFleetspeakTest(flow_test_lib.FlowTestsBaseclass):
  """Test the process listing flow w/ Fleetspeak."""

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

    self.client_id = self.SetupClient(0)
    SetAFF4FSEnabledFlag(self.client_id, token=self.token)

  def testProcessListingOnlyFleetspeak(self):
    """Test that the ListProcesses flow works with Fleetspeak."""
    client_mock = ListProcessesMock([
        rdf_client.Process(
            pid=2,
            ppid=1,
            cmdline=["cmd.exe"],
            exe=r"c:\windows\cmd.exe",
            ctime=1333718907167083L)
    ])

    def SendCallback(fs_msg):
      pb_msg = jobs_pb2.GrrMessage()
      fs_msg.data.Unpack(pb_msg)
      msg = rdf_flows.GrrMessage.FromSerializedString(
          pb_msg.SerializeToString())
      client_mock.mock_task_queue.append(msg)

    service_name = "GRR"
    fake_service_client = _FakeGRPCServiceClient(
        service_name, send_callback=SendCallback)

    fleetspeak_connector.Reset()
    fleetspeak_connector.Init(service_client=fake_service_client)

    with mock.patch.object(
        fake_service_client, "Send", wraps=fake_service_client.Send):
      flow_urn = flow.GRRFlow.StartFlow(
          client_id=self.client_id,
          flow_name=flow_processes.ListProcesses.__name__,
          token=self.token)
      for s in flow_test_lib.TestFlowHelper(
          flow_urn, client_mock, client_id=self.client_id, token=self.token):
        session_id = s

      fleetspeak_connector.CONN.Send.assert_called()

    # Check the output collection
    processes = flow.GRRFlow.ResultCollectionForFID(session_id)

    self.assertEqual(len(processes), 1)
    process, = processes

    self.assertEqual(process.ctime, 1333718907167083L)
    self.assertEqual(process.cmdline, ["cmd.exe"])


def main(args):
  test_lib.main(args)


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