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)
