def testDownloadCollectionWithFoldersEntries(self):
"""Check we can download a collection that also references folders."""
fd = sequential_collection.GeneralIndexedCollection(self.collection_urn)
with data_store.DB.GetMutationPool() as pool:
fd.Add(
rdf_file_finder.FileFinderResult(
stat_entry=rdf_client_fs.StatEntry(
pathspec=rdf_paths.PathSpec(path="testfile5",
pathtype="OS"))),
mutation_pool=pool)
fd.Add(
rdf_file_finder.FileFinderResult(
stat_entry=rdf_client_fs.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.assertIn("testfile5", os.listdir(expected_outdir))
self.assertIn("testdir1", os.listdir(expected_outdir))
def testDownloadCollectionIgnoresArtifactResultsWithoutFiles(self):
# Create a collection with URNs to some files.
fd = sequential_collection.GeneralIndexedCollection(self.collection_urn)
with data_store.DB.GetMutationPool() 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 _VerifyDownload(self):
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("testfile1" in os.listdir(expected_outdir))
self.assertTrue("testfile2" in os.listdir(expected_outdir))
self.assertTrue("testfile5" in os.listdir(expected_outdir))
self.assertTrue("testfile6" in os.listdir(expected_outdir))
# Check we dumped a YAML file to the root of the client.
expected_rootdir = os.path.join(tmpdir, self.client_id.Basename())
self.assertTrue("client_info.yaml" in os.listdir(expected_rootdir))
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)
with data_store.DB.GetMutationPool() as pool:
fd.Add(rdfvalue.RDFURN(self.out.Add("testfile1")), mutation_pool=pool)
fd.Add(
rdf_client_fs.StatEntry(
pathspec=rdf_paths.PathSpec(path="testfile2", pathtype="OS")),
mutation_pool=pool)
fd.Add(
rdf_file_finder.FileFinderResult(
stat_entry=rdf_client_fs.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.assertLen(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"))
