def CheckNotificationsDisappear(self, session_id):
worker_obj = worker.GRRWorker(token=self.token)
manager = queue_manager.QueueManager(token=self.token)
notification = rdf_flows.GrrNotification(session_id=session_id)
with data_store.DB.GetMutationPool() as pool:
manager.NotifyQueue(notification, mutation_pool=pool)
notifications = manager.GetNotificationsByPriority(queues.FLOWS).get(
notification.priority, [])
# Check the notification is there. With multiple worker queue shards we can
# get other notifications such as for audit event listeners, so we need to
# filter out ours.
notifications = [
x for x in notifications if x.session_id == session_id
]
self.assertEqual(len(notifications), 1)
# Process all messages
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
notifications = manager.GetNotificationsByPriority(queues.FLOWS).get(
notification.priority, [])
notifications = [
x for x in notifications if x.session_id == session_id
]
# Check the notification is now gone.
self.assertEqual(len(notifications), 0)
def QueueNotification(self, notification=None, timestamp=None, **kw):
"""Queues a notification for a flow."""
if notification is None:
notification = rdf_flows.GrrNotification(**kw)
session_id = notification.session_id
if session_id:
if timestamp is None:
timestamp = self.frozen_timestamp
# We must not store more than one notification per session id and
# timestamp or there will be race conditions. We therefore only keep
# the one with the highest request number (indicated by last_status).
# Note that timestamp might be None. In that case we also only want
# to keep the latest.
if timestamp is None:
ts_str = "None"
else:
ts_str = int(timestamp)
key = "%s!%s" % (session_id, ts_str)
existing = self.notifications.get(key)
if existing is not None:
if existing[0].last_status < notification.last_status:
self.notifications[key] = (notification, timestamp)
else:
self.notifications[key] = (notification, timestamp)
def testAuthentication2(self):
"""Test that flows refuse to processes unauthenticated messages.
Here we try to simulate an attacker injecting unauthenticated
messages midstream.
The current implementation actually fails to process the entire
flow since the injected messages displace the real ones if they
arrive earlier. This can be an effective DoS against legitimate
clients but would require attackers to guess session ids.
"""
flow_obj = self.FlowSetup("FlowOrderTest")
# Simulate processing messages arriving in random order
message_ids = [1, 2]
self.SendMessages(message_ids, flow_obj.session_id, authenticated=True)
# Now suppose some of the messages are spoofed
message_ids = [3, 4, 5]
self.SendMessages(message_ids, flow_obj.session_id, authenticated=False)
# And now our real messages arrive
message_ids = [5, 6]
self.SendMessages(message_ids, flow_obj.session_id, authenticated=True)
# Send the status message
self.SendOKStatus(7, flow_obj.session_id)
runner = flow_obj.GetRunner()
notification = rdf_flows.GrrNotification(
timestamp=rdfvalue.RDFDatetime.Now())
runner.ProcessCompletedRequests(notification)
# Some messages should actually be processed
self.assertEqual(flow_obj.messages, [1, 2, 5, 6])
def ExpireRules(self):
"""Removes any rules with an expiration date in the past."""
rules = self.Get(self.Schema.RULES)
new_rules = self.Schema.RULES()
now = time.time() * 1e6
expired_session_ids = set()
for rule in rules:
if rule.expires > now:
new_rules.Append(rule)
else:
for action in rule.actions:
if action.hunt_id:
expired_session_ids.add(action.hunt_id)
if expired_session_ids:
# Notify the worker to mark this hunt as terminated.
manager = queue_manager.QueueManager(token=self.token)
manager.MultiNotifyQueue([
rdf_flows.GrrNotification(session_id=session_id)
for session_id in expired_session_ids
])
if len(new_rules) < len(rules):
self.Set(self.Schema.RULES, new_rules)
self.Flush()
def _GetUnsortedNotifications(self,
queue_shard,
notifications_by_session_id=None):
"""Returns all the available notifications for a queue_shard.
Args:
queue_shard: urn of queue shard
notifications_by_session_id: store notifications in this dict rather than
creating a new one
Returns:
dict of notifications. keys are session ids.
"""
if not notifications_by_session_id:
notifications_by_session_id = {}
end_time = self.frozen_timestamp or rdfvalue.RDFDatetime().Now()
for predicate, serialized_notification, ts in data_store.DB.ResolveRegex(
queue_shard,
self.NOTIFY_PREDICATE_PREFIX % ".*",
timestamp=(0, end_time),
token=self.token,
limit=10000):
# Parse the notification.
try:
notification = rdf_flows.GrrNotification(
serialized_notification)
except Exception: # pylint: disable=broad-except
logging.exception(
"Can't unserialize notification, deleting it: "
"predicate=%s, ts=%d", predicate, ts)
data_store.DB.DeleteAttributes(
queue_shard,
[predicate],
token=self.token,
# Make the time range narrow, but be sure to include the needed
# notification.
start=ts,
end=ts,
sync=True)
continue
# Strip the prefix from the predicate to get the session_id.
session_id = predicate[len(self.NOTIFY_PREDICATE_PREFIX % ""):]
notification.session_id = session_id
notification.timestamp = ts
existing = notifications_by_session_id.get(notification.session_id)
if existing:
# If we have a notification for this session_id already, we only store
# the one that was scheduled last.
if notification.first_queued > existing.first_queued:
notifications_by_session_id[
notification.session_id] = notification
else:
notifications_by_session_id[
notification.session_id] = notification
return notifications_by_session_id
def QueueNotification(self, notification=None, timestamp=None, **kw):
"""Queues a notification for a flow."""
if notification is None:
notification = rdf_flows.GrrNotification(**kw)
if notification.session_id:
if timestamp is None:
timestamp = self.frozen_timestamp
self.notifications.append((notification, timestamp))
def testNotificationReschedulingTTL(self):
"""Test that notifications are not rescheduled forever."""
with test_lib.FakeTime(10000):
worker_obj = worker.GRRWorker(token=self.token)
flow_obj = self.FlowSetup("RaisingTestFlow")
session_id = flow_obj.session_id
flow_obj.Close()
with queue_manager.QueueManager(token=self.token) as manager:
notification = rdf_flows.GrrNotification(session_id=session_id,
timestamp=time.time(),
last_status=1)
with data_store.DB.GetMutationPool() as pool:
manager.NotifyQueue(notification, mutation_pool=pool)
notifications = manager.GetNotifications(queues.FLOWS)
# Check the notification is there.
notifications = [
n for n in notifications if n.session_id == session_id
]
self.assertEqual(len(notifications), 1)
delay = flow_runner.FlowRunner.notification_retry_interval
ttl = notification.ttl
for i in xrange(ttl - 1):
with test_lib.FakeTime(10000 + 100 + delay * (i + 1)):
# Process all messages.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
notifications = manager.GetNotifications(queues.FLOWS)
# Check the notification is for the correct session_id.
notifications = [
n for n in notifications if n.session_id == session_id
]
self.assertEqual(len(notifications), 1)
with test_lib.FakeTime(10000 + 100 + delay * ttl):
# Process all messages.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
notifications = manager.GetNotifications(queues.FLOWS)
self.assertEqual(len(notifications), 0)
def testReordering(self):
"""Check that out of order client messages are reordered."""
flow_obj = self.FlowSetup(flow_test_lib.FlowOrderTest.__name__)
# Simulate 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
self.SendOKStatus(6, flow_obj.session_id)
runner = flow_obj.GetRunner()
notification = rdf_flows.GrrNotification(
timestamp=rdfvalue.RDFDatetime.Now())
runner.ProcessCompletedRequests(notification)
# Check that the messages were processed in order
self.assertEqual(flow_obj.messages, [1, 2, 3, 4, 5])
def testAuthentication1(self):
"""Test that flows refuse to processes unauthenticated messages."""
flow_obj = self.FlowSetup("FlowOrderTest")
# Simulate processing messages arriving in random order
message_ids = [2, 1, 4, 3, 5]
self.SendMessages(message_ids, flow_obj.session_id, authenticated=False)
# Send the status message
self.SendOKStatus(6, flow_obj.session_id)
runner = flow_obj.GetRunner()
notification = rdf_flows.GrrNotification(
timestamp=rdfvalue.RDFDatetime.Now())
runner.ProcessCompletedRequests(notification)
# Now messages should actually be processed
self.assertEqual(flow_obj.messages, [])
def testNotificationRacesAreResolved(self):
# We need a random flow object for this test.
session_id = flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
worker_obj = worker.GRRWorker(token=self.token)
manager = queue_manager.QueueManager(token=self.token)
manager.DeleteNotification(session_id)
manager.Flush()
# We simulate a race condition here - the notification for request #1 is
# there but the actual request #1 is not. The worker should pick up the
# notification, notice that the request #1 is not there yet and reschedule
# the notification.
notification = rdf_flows.GrrNotification(session_id=session_id,
last_status=1)
with data_store.DB.GetMutationPool() as pool:
manager.NotifyQueue(notification, mutation_pool=pool)
notifications = manager.GetNotifications(queues.FLOWS)
# Check the notification is there.
notifications = [
n for n in notifications if n.session_id == session_id
]
self.assertEqual(len(notifications), 1)
# Process all messages
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
delay = flow_runner.FlowRunner.notification_retry_interval
with test_lib.FakeTime(time.time() + 10 + delay):
requeued_notifications = manager.GetNotifications(queues.FLOWS)
# Check that there is a new notification.
notifications = [
n for n in notifications if n.session_id == session_id
]
self.assertEqual(len(requeued_notifications), 1)
self.assertEqual(requeued_notifications[0].first_queued,
notifications[0].first_queued)
self.assertNotEqual(requeued_notifications[0].timestamp,
notifications[0].timestamp)
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))