def testReceiveMessageList(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) 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()) fs_message.validation_info.tags["foo"] = "bar" 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)) self.assertEqual( client_data[client_id].last_fleetspeak_validation_info. ToStringDict(), {"foo": "bar"}) 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 main(argv): del argv # Unused. config.CONFIG.AddContext("FleetspeakFrontend Context") server_startup.Init() server_startup.DropPrivileges() fleetspeak_connector.Init() fsd = fleetspeak_frontend_server.GRRFSServer() fleetspeak_connector.CONN.Listen(fsd.Process) logging.info("Serving through Fleetspeak ...") try: while True: time.sleep(600) except KeyboardInterrupt: print("Caught keyboard interrupt, stopping")
def testWriteLastPingForNewClients(self): fs_server = fleetspeak_frontend_server.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)