def testReceiveMessageListFleetspeak(self): 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 range(1, num_msgs + 1) ] fs_client_id = b"\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=FS_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)
def testReceiveMessageList_Relational(self): if not data_store.RelationalDBEnabled(): self.skipTest("Rel-db-only test.") fs_server = fs_frontend_tool.GRRFSServer() client_id = "C.1234567890123456" flow_id = "12345678" data_store.REL_DB.WriteClientMetadata(client_id, fleetspeak_enabled=True) rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow_id, create_time=rdfvalue.RDFDatetime.Now()) data_store.REL_DB.WriteFlowObject(rdf_flow) flow_request = rdf_flow_objects.FlowRequest(client_id=client_id, flow_id=flow_id, request_id=1) data_store.REL_DB.WriteFlowRequests([flow_request]) session_id = "%s/%s" % (client_id, flow_id) fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id) grr_messages = [] for i in range(1, 10): grr_message = rdf_flows.GrrMessage(request_id=1, response_id=i + 1, session_id=session_id, payload=rdfvalue.RDFInteger(i)) grr_messages.append(grr_message) packed_messages = rdf_flows.PackedMessageList() communicator.Communicator.EncodeMessageList( rdf_flows.MessageList(job=grr_messages), packed_messages) fs_message = fs_common_pb2.Message(message_type="MessageList", source=fs_common_pb2.Address( client_id=fs_client_id, service_name=FS_SERVICE_NAME)) fs_message.data.Pack(packed_messages.AsPrimitiveProto()) with test_lib.FakeTime( rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)): fs_server.Process(fs_message, None) # Ensure the last-ping timestamp gets updated. client_data = data_store.REL_DB.MultiReadClientMetadata([client_id]) self.assertEqual(client_data[client_id].ping, rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)) flow_data = data_store.REL_DB.ReadAllFlowRequestsAndResponses( client_id, flow_id) self.assertLen(flow_data, 1) stored_flow_request, flow_responses = flow_data[0] self.assertEqual(stored_flow_request, flow_request) self.assertLen(flow_responses, 9)
def testWriteLastPingForNewClients(self): if not data_store.RelationalDBEnabled(): self.skipTest("Rel-db-only test.") fs_server = fs_frontend_tool.GRRFSServer() client_id = "C.1234567890123456" flow_id = "12345678" session_id = "%s/%s" % (client_id, flow_id) fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id) grr_message = rdf_flows.GrrMessage(request_id=1, response_id=1, session_id=session_id, payload=rdfvalue.RDFInteger(1)) fs_message = fs_common_pb2.Message(message_type="GrrMessage", source=fs_common_pb2.Address( client_id=fs_client_id, service_name=FS_SERVICE_NAME)) fs_message.data.Pack(grr_message.AsPrimitiveProto()) fake_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123) with mock.patch.object( events.Events, "PublishEvent", wraps=events.Events.PublishEvent) as publish_event_fn: with mock.patch.object(data_store.REL_DB, "WriteClientMetadata", wraps=data_store.REL_DB.WriteClientMetadata ) as write_metadata_fn: with test_lib.FakeTime(fake_time): fs_server.Process(fs_message, None) self.assertEqual(write_metadata_fn.call_count, 1) client_data = data_store.REL_DB.MultiReadClientMetadata( [client_id]) self.assertEqual(client_data[client_id].ping, fake_time) # TODO(user): publish_event_fn.assert_any_call( # "ClientEnrollment", mock.ANY, token=mock.ANY) doesn't work here # for some reason. triggered_events = [] for call_args, _ in publish_event_fn.call_args_list: if call_args: triggered_events.append(call_args[0]) self.assertIn("ClientEnrollment", triggered_events)