def testDownloadHuntResultCollection(self):
"""Check we can download files references in HuntResultCollection."""
# Create a collection with URNs to some files.
fd = results.HuntResultCollection(self.collection_urn,
token=self.token)
with data_store.DB.GetMutationPool(token=self.token) as pool:
fd.AddAsMessage(rdfvalue.RDFURN(self.out.Add("testfile1")),
self.client_id,
mutation_pool=pool)
fd.AddAsMessage(rdf_client.StatEntry(
pathspec=rdf_paths.PathSpec(path="testfile2", pathtype="OS")),
self.client_id,
mutation_pool=pool)
fd.AddAsMessage(rdf_file_finder.FileFinderResult(
stat_entry=rdf_client.StatEntry(pathspec=rdf_paths.PathSpec(
path="testfile5", pathtype="OS"))),
self.client_id,
mutation_pool=pool)
fd.AddAsMessage(collectors.ArtifactFilesDownloaderResult(
downloaded_file=rdf_client.StatEntry(
pathspec=rdf_paths.PathSpec(path="testfile6",
pathtype="OS"))),
self.client_id,
mutation_pool=pool)
self._VerifyDownload()
def _OpenCollectionPath(coll_path):
"""Tries to open various types of collections at the given path."""
collection = results.HuntResultCollection(coll_path)
if collection and collection[0].payload:
return collection
collection = sequential_collection.GeneralIndexedCollection(coll_path)
if collection:
return collection
def testEmptyQueue(self):
# Create and empty HuntResultCollection.
collection_urn = rdfvalue.RDFURN("aff4:/testEmptyQueue/collection")
hunts_results.HuntResultCollection(collection_urn)
# The queue starts empty, and returns no notifications.
results = hunts_results.HuntResultQueue.ClaimNotificationsForCollection(
token=self.token)
self.assertEqual(None, results[0])
self.assertEqual([], results[1])
def testNotificationsSplitByCollection(self):
# Create two HuntResultCollections.
collection_urn_1 = rdfvalue.RDFURN(
"aff4:/testNotificationsSplitByCollection/collection_1")
collection_urn_2 = rdfvalue.RDFURN(
"aff4:/testNotificationsSplitByCollection/collection_2")
# Add 100 records to each collection, in an interleaved manner.
with data_store.DB.GetMutationPool(token=self.token) as pool:
for i in range(100):
hunts_results.HuntResultCollection.StaticAdd(
collection_urn_1,
rdf_flows.GrrMessage(request_id=i),
mutation_pool=pool)
hunts_results.HuntResultCollection.StaticAdd(
collection_urn_2,
rdf_flows.GrrMessage(request_id=100 + i),
mutation_pool=pool)
# The first result was added to collection 1, so this should return
# all 100 results for collection 1.
results_1 = hunts_results.HuntResultQueue.ClaimNotificationsForCollection(
token=self.token)
self.assertEqual(collection_urn_1, results_1[0])
self.assertEqual(100, len(results_1[1]))
# The first call claimed all the notifications for collection 1. These are
# claimed, so another call should skip them and give all notifications for
# collection 2.
results_2 = hunts_results.HuntResultQueue.ClaimNotificationsForCollection(
token=self.token)
self.assertEqual(collection_urn_2, results_2[0])
self.assertEqual(100, len(results_2[1]))
values_read = []
collection_2 = hunts_results.HuntResultCollection(
collection_urn_2, token=self.token)
for message in collection_2.MultiResolve([(ts, suffix)
for (_, ts,
suffix) in results_2[1]]):
values_read.append(message.request_id)
self.assertEqual(sorted(values_read), range(100, 200))
def testNotificationsContainTimestamps(self):
collection_urn = rdfvalue.RDFURN(
"aff4:/testNotificationsContainTimestamps/collection")
with data_store.DB.GetMutationPool() as pool:
for i in range(5):
hunts_results.HuntResultCollection.StaticAdd(
collection_urn,
rdf_flows.GrrMessage(request_id=i),
mutation_pool=pool)
# If we claim results, we should get all 5.
results = hunts_results.HuntResultQueue.ClaimNotificationsForCollection(
token=self.token)
self.assertEqual(collection_urn, results[0])
self.assertEqual(5, len(results[1]))
# Read all the results, using the contained (ts, suffix) pairs.
values_read = []
collection = hunts_results.HuntResultCollection(collection_urn)
for message in collection.MultiResolve(
[r.value.ResultRecord() for r in results[1]]):
values_read.append(message.request_id)
self.assertEqual(sorted(values_read), range(5))
