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([
rdfvalue.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 = rdfvalue.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 = rdfvalue.GrrNotification(**kw)
if notification.session_id:
if timestamp is None:
timestamp = self.frozen_timestamp
self.notifications.append((notification, timestamp))
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(worker.DEFAULT_WORKER_QUEUE,
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 = rdfvalue.GrrNotification(session_id=session_id,
last_status=1)
manager.NotifyQueue(notification)
notifications = manager.GetNotifications(worker.DEFAULT_WORKER_QUEUE)
# 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 = config_lib.CONFIG["Worker.notification_retry_interval"]
with test_lib.FakeTime(time.time() + 10 + delay):
requeued_notifications = manager.GetNotifications(
worker.DEFAULT_WORKER_QUEUE)
# 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 CheckNotificationsDisappear(self, session_id):
worker_obj = worker.GRRWorker(worker.DEFAULT_WORKER_QUEUE,
token=self.token)
manager = queue_manager.QueueManager(token=self.token)
notification = rdfvalue.GrrNotification(session_id=session_id)
manager.NotifyQueue(notification)
notifications = manager.GetNotificationsByPriority(
worker.DEFAULT_WORKER_QUEUE).get(notification.priority, [])
# Check the notification is there.
self.assertEqual(len(notifications), 1)
self.assertEqual(notifications[0].session_id, session_id)
# Process all messages
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
notifications = manager.GetNotificationsByPriority(
worker.DEFAULT_WORKER_QUEUE).get(notification.priority, [])
# Check the notification is now gone.
self.assertEqual(len(notifications), 0)