def Notify(self, message_type, subject, msg):
"""Send a notification to the originating user.
Args:
message_type: The type of the message. This allows the UI to format
a link to the original object e.g. "ViewObject" or "HostInformation"
subject: The urn of the AFF4 object of interest in this link.
msg: A free form textual message.
"""
user = self.context.creator
# Don't send notifications to system users.
if (self.runner_args.notify_to_user
and user not in aff4_users.GRRUser.SYSTEM_USERS):
# Prefix the message with the hostname of the client we are running
# against.
if self.runner_args.client_id:
client_fd = aff4.FACTORY.Open(self.runner_args.client_id,
mode="rw",
token=self.token)
hostname = client_fd.Get(client_fd.Schema.HOSTNAME) or ""
client_msg = "%s: %s" % (hostname, msg)
else:
client_msg = msg
# Add notification to the User object.
with aff4.FACTORY.Create(aff4.ROOT_URN.Add("users").Add(user),
aff4_users.GRRUser,
mode="rw",
token=self.token) as fd:
# Queue notifications to the user.
fd.Notify(message_type, subject, client_msg, self.session_id)
# Add notifications to the flow.
notification = rdf_flows.Notification(
type=message_type,
subject=utils.SmartUnicode(subject),
message=utils.SmartUnicode(msg),
source=self.session_id,
timestamp=rdfvalue.RDFDatetime.Now())
# TODO(amoser): This should go into the DB api.
data_store.DB.Set(self.session_id,
self.flow_obj.Schema.NOTIFICATION,
notification,
replace=False,
sync=True)
# Disable further notifications.
self.context.user_notified = True
def testReordering(self):
"""Check that out of order client messages are reordered."""
flow_obj = self.FlowSetup("FlowOrderTest")
# Simultate processing messages arriving in random order
message_ids = [2, 1, 4, 3, 5]
self.SendMessages(message_ids, flow_obj.session_id)
# Send the status message
message = self.SendOKStatus(6, flow_obj.session_id)
runner = flow_runner.FlowRunner(flow_obj)
notification = rdf_flows.Notification(
timestamp=rdfvalue.RDFDatetime().Now())
runner.ProcessCompletedRequests(notification, [message])
# Check that the messages were processed in order
self.assertEqual(flow_obj.messages, [1, 2, 3, 4, 5])
