def testDownloadCollectionWithFoldersEntries(self):
"""Check we can download a collection that also references folders."""
fd = sequential_collection.GeneralIndexedCollection(
self.collection_urn, token=self.token)
with data_store.DB.GetMutationPool(token=self.token) as pool:
fd.Add(rdf_file_finder.FileFinderResult(
stat_entry=rdf_client.StatEntry(pathspec=rdf_paths.PathSpec(
path="testfile5", pathtype="OS"))),
mutation_pool=pool)
fd.Add(rdf_file_finder.FileFinderResult(
stat_entry=rdf_client.StatEntry(pathspec=rdf_paths.PathSpec(
path="testdir1", pathtype="OS"),
st_mode=stat.S_IFDIR)),
mutation_pool=pool)
with utils.TempDirectory() as tmpdir:
export_utils.DownloadCollection(self.collection_urn,
tmpdir,
overwrite=True,
dump_client_info=True,
token=self.token,
max_threads=2)
expected_outdir = os.path.join(tmpdir, self.out.Path()[1:])
# Check we found both files.
self.assertTrue("testfile5" in os.listdir(expected_outdir))
self.assertTrue("testdir1" in os.listdir(expected_outdir))
def ResultCollectionForArtifact(cls,
session_id,
artifact_name,
token=None):
urn = rdfvalue.RDFURN("_".join(
(str(session_id.Add(flow.RESULTS_SUFFIX)),
utils.SmartStr(artifact_name))))
return sequential_collection.GeneralIndexedCollection(urn)
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 setUp(self):
super(FilterCollectionTest, self).setUp()
self.fd = sequential_collection.GeneralIndexedCollection(
rdfvalue.RDFURN("aff4:/tmp/foo/bar"))
with data_store.DB.GetMutationPool() as pool:
for i in range(10):
self.fd.Add(
rdf_paths.PathSpec(path="/var/os/tmp-%d" % i, pathtype="OS"),
mutation_pool=pool)
def ResultCollectionForFID(cls, flow_id):
"""Returns the ResultCollection for the flow with a given flow_id.
Args:
flow_id: The id of the flow, a RDFURN of the form aff4:/flows/F:123456.
Returns:
The collection containing the results for the flow identified by the id.
"""
return sequential_collection.GeneralIndexedCollection(
flow_id.Add(RESULTS_SUFFIX))
def testAddGet(self):
collection = sequential_collection.GeneralIndexedCollection(
rdfvalue.RDFURN("aff4:/sequential_collection/testAddGetIndexed"))
with data_store.DB.GetMutationPool() as pool:
collection.Add(rdfvalue.RDFInteger(42), mutation_pool=pool)
collection.Add(rdfvalue.RDFString("the meaning of life"),
mutation_pool=pool)
self.assertEqual(collection[0].__class__, rdfvalue.RDFInteger)
self.assertEqual(collection[0], 42)
self.assertEqual(collection[1].__class__, rdfvalue.RDFString)
self.assertEqual(collection[1], "the meaning of life")
def testDownloadCollectionIgnoresArtifactResultsWithoutFiles(self):
# Create a collection with URNs to some files.
fd = sequential_collection.GeneralIndexedCollection(
self.collection_urn, token=self.token)
with data_store.DB.GetMutationPool(token=self.token) as pool:
fd.Add(collectors.ArtifactFilesDownloaderResult(),
mutation_pool=pool)
with utils.TempDirectory() as tmpdir:
export_utils.DownloadCollection(self.collection_urn,
tmpdir,
overwrite=True,
dump_client_info=True,
token=self.token,
max_threads=2)
expected_outdir = os.path.join(tmpdir, self.out.Path()[1:])
self.assertFalse(os.path.exists(expected_outdir))
def testDownloadCollectionWithFlattenOption(self):
"""Check we can download files references in a collection."""
# Create a collection with URNs to some files.
fd = sequential_collection.GeneralIndexedCollection(
self.collection_urn, token=self.token)
with data_store.DB.GetMutationPool(token=self.token) as pool:
fd.Add(rdfvalue.RDFURN(self.out.Add("testfile1")),
mutation_pool=pool)
fd.Add(rdf_client.StatEntry(
pathspec=rdf_paths.PathSpec(path="testfile2", pathtype="OS")),
mutation_pool=pool)
fd.Add(rdf_file_finder.FileFinderResult(
stat_entry=rdf_client.StatEntry(pathspec=rdf_paths.PathSpec(
path="testfile5", pathtype="OS"))),
mutation_pool=pool)
with utils.TempDirectory() as tmpdir:
export_utils.DownloadCollection(self.collection_urn,
tmpdir,
overwrite=True,
dump_client_info=True,
flatten=True,
token=self.token,
max_threads=2)
# Check that "files" folder is filled with symlinks to downloaded files.
symlinks = os.listdir(os.path.join(tmpdir, "files"))
self.assertEqual(len(symlinks), 3)
self.assertListEqual(sorted(symlinks), [
"C.1000000000000000_fs_os_testfile1",
"C.1000000000000000_fs_os_testfile2",
"C.1000000000000000_fs_os_testfile5"
])
self.assertEqual(
os.readlink(
os.path.join(tmpdir, "files",
"C.1000000000000000_fs_os_testfile1")),
os.path.join(tmpdir, "C.1000000000000000", "fs", "os",
"testfile1"))
def testDownloadGeneralIndexedCollection(self):
"""Check we can download files references in GeneralIndexedCollection."""
fd = sequential_collection.GeneralIndexedCollection(
self.collection_urn, token=self.token)
self._AddTestData(fd)
self._VerifyDownload()
