def FetchAndRemoveRequestsAndResponses(self, session_id):
"""Removes WellKnownFlow messages from the queue and returns them."""
messages = []
with queue_manager.WellKnownQueueManager(token=self.token) as manager:
for _, responses in manager.FetchRequestsAndResponses(session_id):
messages.extend(responses)
manager.DeleteWellKnownFlowResponses(session_id, responses)
return messages
def testWellKnownFlowsNotifications(self):
test_lib.WellKnownSessionTest.messages = []
test_lib.WellKnownSessionTest2.messages = []
session_id1 = test_lib.WellKnownSessionTest.well_known_session_id
session_id2 = test_lib.WellKnownSessionTest2.well_known_session_id
messages = []
for i in range(1, 5):
messages.append(
rdf_flows.GrrMessage(request_id=0,
response_id=0,
session_id=session_id1,
payload=rdfvalue.RDFInteger(i)))
messages.append(
rdf_flows.GrrMessage(request_id=0,
response_id=0,
session_id=session_id2,
payload=rdfvalue.RDFInteger(i)))
# This test whitelists only one flow.
self.assertIn(session_id1.FlowName(), self.server.well_known_flows)
self.assertNotIn(session_id2.FlowName(), self.server.well_known_flows)
self.server.ReceiveMessages(self.client_id, messages)
# Wait for async actions to complete
self.server.thread_pool.Join()
# Flow 1 should have been processed right away.
test_lib.WellKnownSessionTest.messages.sort()
self.assertEqual(test_lib.WellKnownSessionTest.messages,
list(range(1, 5)))
# But not Flow 2.
self.assertEqual(test_lib.WellKnownSessionTest2.messages, [])
manager = queue_manager.WellKnownQueueManager(token=self.token)
notifications = manager.GetNotificationsForAllShards(
session_id1.Queue())
# Flow 1 was proecessed on the frontend, no queued responses available.
responses = list(manager.FetchRequestsAndResponses(session_id1))
self.assertEqual(responses, [])
# And also no notifications.
self.assertNotIn(
session_id1,
[notification.session_id for notification in notifications])
# But for Flow 2 there should be some responses + a notification.
responses = list(manager.FetchRequestsAndResponses(session_id2))
self.assertEqual(len(responses), 4)
self.assertIn(
session_id2,
[notification.session_id for notification in notifications])
def PublishMultipleEvents(cls, events, delay=None, token=None):
"""Publish the message into all listeners of the event.
If event_name is a string, we send the message to all event handlers which
contain this string in their EVENT static member. This allows the event to
be sent to multiple interested listeners. Alternatively, the event_name can
specify a single URN of an event listener to receive the message.
Args:
events: A dict with keys being event names and values being lists of
messages.
delay: An rdfvalue.Duration object. If given, the event will be published
after the indicated time.
token: ACL token.
Raises:
ValueError: If the message is invalid. The message must be a Semantic
Value (instance of RDFValue) or a full GrrMessage.
"""
with queue_manager.WellKnownQueueManager(token=token) as manager:
event_name_map = registry.EventRegistry.EVENT_NAME_MAP
for event_name, messages in events.iteritems():
handler_urns = []
if isinstance(event_name, basestring):
for event_cls in event_name_map.get(event_name, []):
if event_cls.well_known_session_id is None:
logging.error(
"Well known flow %s has no session_id.",
event_cls.__name__)
else:
handler_urns.append(
event_cls.well_known_session_id)
else:
handler_urns.append(event_name)
for msg in messages:
# Allow the event name to be either a string or a URN of an event
# listener.
if not isinstance(msg, rdfvalue.RDFValue):
raise ValueError(
"Can only publish RDFValue instances.")
# Wrap the message in a GrrMessage if needed.
if not isinstance(msg, rdf_flows.GrrMessage):
msg = rdf_flows.GrrMessage(payload=msg)
# Randomize the response id or events will get overwritten.
msg.response_id = msg.task_id = msg.GenerateTaskID()
# Well known flows always listen for request id 0.
msg.request_id = 0
timestamp = None
if delay:
timestamp = (rdfvalue.RDFDatetime.Now() +
delay).AsMicroSecondsFromEpoch()
# Forward the message to the well known flow's queue.
for event_urn in handler_urns:
manager.QueueResponse(event_urn, msg)
manager.QueueNotification(
rdf_flows.GrrNotification(session_id=event_urn,
priority=msg.priority,
timestamp=timestamp))