def SendGrrMessageThroughFleetspeak(grr_id, grr_msg): """Sends the given GrrMessage through FS.""" fs_msg = fs_common_pb2.Message(message_type="GrrMessage", destination=fs_common_pb2.Address( client_id=GRRIDToFleetspeakID(grr_id), service_name="GRR")) fs_msg.data.Pack(grr_msg.AsPrimitiveProto()) if grr_msg.session_id is not None: annotation = fs_msg.annotations.entries.add() annotation.key, annotation.value = "flow_id", grr_msg.session_id.Basename( ) if grr_msg.request_id is not None: annotation = fs_msg.annotations.entries.add() annotation.key, annotation.value = "request_id", str( grr_msg.request_id) fleetspeak_connector.CONN.outgoing.InsertMessage(fs_msg)
def testReceiveMessages(self): fs_server = fleetspeak_frontend_server.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) fs_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)) 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()) fs_messages.append(fs_message) with test_lib.FakeTime(rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)): for fs_message in fs_messages: 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 testPingIsRecorded(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 = self.SetupTestClientObject(grr_client_nr) self.SetupClient(grr_client_nr) messages = [ rdf_flows.GrrMessage(request_id=1, response_id=1, session_id="F:123456", payload=rdfvalue.RDFInteger(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.client_id)) 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()) fake_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(42) with test_lib.FakeTime(fake_time): fsd.Process(fs_message, None) md = data_store.REL_DB.ReadClientMetadata(grr_client.client_id) self.assertEqual(md.ping, fake_time) with aff4.FACTORY.Open(grr_client.client_id) as client: self.assertEqual(client.Get(client.Schema.PING), fake_time)
def _SendMessages(self, grr_msgs, background=False): """Sends a block of messages through Fleetspeak.""" message_list = rdf_flows.PackedMessageList() communicator.Communicator.EncodeMessageList( rdf_flows.MessageList(job=grr_msgs), message_list) fs_msg = fs_common_pb2.Message( message_type="MessageList", destination=fs_common_pb2.Address(service_name="GRR"), background=background) fs_msg.data.Pack(message_list.AsPrimitiveProto()) try: sent_bytes = self._fs.Send(fs_msg) except (IOError, struct.error) as e: logging.critical("Broken local Fleetspeak connection (write end).") raise e stats.STATS.IncrementCounter("grr_client_sent_bytes", sent_bytes)
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)
def _SendMessages(self, grr_msgs, background=False): """Sends a block of messages through Fleetspeak.""" message_list = rdf_flows.PackedMessageList() communicator.Communicator.EncodeMessageList( rdf_flows.MessageList(job=grr_msgs), message_list) fs_msg = fs_common_pb2.Message( message_type="MessageList", destination=fs_common_pb2.Address(service_name="GRR"), background=background) fs_msg.data.Pack(message_list.AsPrimitiveProto()) try: sent_bytes = self._fs.Send(fs_msg) except (IOError, struct.error) as e: logging.fatal("Broken local Fleetspeak connection (write end): %r", e, exc_info=True) # The fatal call above doesn't terminate the program. The reasons for # this might include Python threads persistency, or Python logging # mechanisms' inconsistency. os._exit(1) # pylint: disable=protected-access stats.STATS.IncrementCounter("grr_client_sent_bytes", sent_bytes)
def testReceiveMessagesFleetspeak(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())) fs_messages = [ fs_common_pb2.Message(message_type="GrrMessage", source=fs_common_pb2.Address( client_id=fs_client_id, service_name=FS_SERVICE_NAME)) for _ in range(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)