def testSendMessagesWithAnnotations(self, mock_worker_class, mock_conn_class): # We stub out the worker class since it starts threads in its # __init__ method. del mock_worker_class # Unused mock_conn = mock.Mock() mock_conn.Send.return_value = 123 mock_conn_class.return_value = mock_conn client_id = "C.0123456789abcdef" flow_id = "01234567" client = fleetspeak_client.GRRFleetspeakClient() grr_messages = [] expected_annotations = fs_common_pb2.Annotations() # 500 bytes translates to ~19 annotations. while expected_annotations.ByteSize() < 500: grr_message = rdf_flows.GrrMessage( session_id="%s/%s" % (client_id, flow_id), name="TestClientAction", request_id=2, response_id=len(grr_messages) + 1) annotation = expected_annotations.entries.add() annotation.key = fleetspeak_client._DATA_IDS_ANNOTATION_KEY annotation.value = "%s:2:%d" % (flow_id, len(grr_messages) + 1) grr_messages.append(grr_message) client._sender_queue.put(grr_message) # Add an extra GrrMessage whose annotation will not be captured. extra_message = rdf_flows.GrrMessage(session_id="%s/%s" % (client_id, flow_id), name="TestClientAction", request_id=3, response_id=1) grr_messages.append(extra_message) client._sender_queue.put(extra_message) self.assertLess(len(grr_messages), fleetspeak_client._MAX_MSG_LIST_MSG_COUNT) self.assertLess(sum(len(x.SerializeToBytes()) for x in grr_messages), fleetspeak_client._MAX_MSG_LIST_BYTES) client._SendOp() mock_conn.Send.assert_called_once() send_args, _ = mock_conn.Send.call_args fs_message = send_args[0] packed_message_list = rdf_flows.PackedMessageList.protobuf() fs_message.data.Unpack(packed_message_list) message_list = communicator.Communicator.DecompressMessageList( rdf_flows.PackedMessageList.FromSerializedBytes( packed_message_list.SerializeToString())) self.assertListEqual(list(message_list.job), grr_messages) self.assertEqual(fs_message.annotations, expected_annotations)
def testSendGrrMessage(self, mock_conn): client_id = "C.0123456789abcdef" flow_id = "01234567" grr_message = rdf_flows.GrrMessage( session_id="%s/%s" % (client_id, flow_id), name="TestClientAction", request_id=1) fleetspeak_utils.SendGrrMessageThroughFleetspeak(client_id, grr_message) mock_conn.outgoing.InsertMessage.assert_called_once() insert_args, _ = mock_conn.outgoing.InsertMessage.call_args fs_message = insert_args[0] expected_annotations = common_pb2.Annotations(entries=[ common_pb2.Annotations.Entry(key="flow_id", value=flow_id), common_pb2.Annotations.Entry(key="request_id", value="1"), ]) unpacked_message = rdf_flows.GrrMessage.protobuf() fs_message.data.Unpack(unpacked_message) self.assertEqual(fs_message.annotations, expected_annotations) self.assertEqual(grr_message.AsPrimitiveProto(), unpacked_message)