def testMissingFilesAreListedInManifest(self):
generator = archive_generator.FlowArchiveGenerator(
self.flow, archive_generator.ArchiveFormat.ZIP)
mappings = [
flow_base.ClientPathArchiveMapping(self.path1, "foo/file"),
flow_base.ClientPathArchiveMapping(
db.ClientPath.OS(self.client_id, ["non", "existing"]),
"foo/bar/file"),
]
fd_path = self._GenerateArchive(generator.Generate(mappings))
zip_fd = zipfile.ZipFile(fd_path)
names = [str(s) for s in sorted(zip_fd.namelist())]
# Expecting in the archive: 1 file (the other shouldn't be found)
# and a manifest.
self.assertLen(names, 2)
contents = zip_fd.read(os.path.join(generator.prefix, "foo", "file"))
self.assertEqual(contents, b"hello1")
manifest = yaml.safe_load(
zip_fd.read(os.path.join(generator.prefix, "MANIFEST")))
self.assertCountEqual(manifest["processed_files"].items(),
[(self.path1.vfs_path, "foo/file")])
self.assertCountEqual(manifest["missing_files"],
["fs/os/non/existing"])
def testCreatesTarContainingTwoMappedFilesAndManifest(self):
generator = archive_generator.FlowArchiveGenerator(
self.flow, archive_generator.ArchiveFormat.TAR_GZ)
mappings = [
flow_base.ClientPathArchiveMapping(self.path1, "foo/file"),
flow_base.ClientPathArchiveMapping(self.path2, "foo/bar/file"),
]
fd_path = self._GenerateArchive(generator.Generate(mappings))
with tarfile.open(fd_path, encoding="utf-8") as tar_fd:
self.assertLen(tar_fd.getnames(), 3)
contents = tar_fd.extractfile(
os.path.join(generator.prefix, "foo", "file")).read()
self.assertEqual(contents, b"hello1")
contents = tar_fd.extractfile(
os.path.join(generator.prefix, "foo", "bar", "file")).read()
self.assertEqual(contents, b"hello2")
manifest = yaml.safe_load(
tar_fd.extractfile(os.path.join(generator.prefix,
"MANIFEST")).read())
self.assertCountEqual(manifest["processed_files"].items(),
[(self.path1.vfs_path, "foo/file"),
(self.path2.vfs_path, "foo/bar/file")])
self.assertCountEqual(manifest["missing_files"], [])
self.assertEqual(manifest["client_id"], self.client_id)
self.assertEqual(manifest["flow_id"], self.flow_id)
def testCreatesZipContainingTwoMappedFilesAndManifest(self):
generator = archive_generator.FlowArchiveGenerator(
self.flow, archive_generator.ArchiveFormat.ZIP)
mappings = [
flow_base.ClientPathArchiveMapping(self.path1, "foo/file"),
flow_base.ClientPathArchiveMapping(self.path2, "foo/bar/file"),
]
fd_path = self._GenerateArchive(generator.Generate(mappings))
zip_fd = zipfile.ZipFile(fd_path)
names = [str(s) for s in sorted(zip_fd.namelist())]
# Expecting in the archive: 2 files and a manifest.
self.assertLen(names, 3)
contents = zip_fd.read(os.path.join(generator.prefix, "foo", "file"))
self.assertEqual(contents, b"hello1")
contents = zip_fd.read(
os.path.join(generator.prefix, "foo", "bar", "file"))
self.assertEqual(contents, b"hello2")
manifest = yaml.safe_load(
zip_fd.read(os.path.join(generator.prefix, "MANIFEST")))
self.assertCountEqual(manifest["processed_files"].items(),
[(self.path1.vfs_path, "foo/file"),
(self.path2.vfs_path, "foo/bar/file")])
self.assertCountEqual(manifest["missing_files"], [])
self.assertEqual(manifest["client_id"], self.client_id)
self.assertEqual(manifest["flow_id"], self.flow_id)
def testPropagatesStreamingExceptions(self):
generator = archive_generator.FlowArchiveGenerator(
self.flow, archive_generator.ArchiveFormat.TAR_GZ)
mappings = [
flow_base.ClientPathArchiveMapping(self.path1, "foo/file"),
flow_base.ClientPathArchiveMapping(self.path2, "foo/bar/file"),
]
with mock.patch.object(
file_store, "StreamFilesChunks", side_effect=Exception("foobar")):
with self.assertRaises(Exception) as context:
self._GenerateArchive(generator.Generate(mappings))
self.assertEqual(str(context.exception), "foobar")
def Handle(self, args, context=None):
flow_object, flow_results = self._GetFlow(args, context)
flow_api_object = ApiFlow().InitFromFlowObject(flow_object)
flow_instance = flow_base.FlowBase.CreateFlowInstance(flow_object)
try:
mappings = flow_instance.GetFilesArchiveMappings(flow_results)
except NotImplementedError:
mappings = None
description = ("Files downloaded by flow %s (%s) that ran on client %s by "
"user %s on %s" %
(flow_api_object.name, args.flow_id, args.client_id,
flow_api_object.creator, flow_api_object.started_at))
target_file_prefix = "%s_flow_%s_%s" % (
args.client_id, flow_api_object.name, str(
flow_api_object.flow_id).replace(":", "_"))
if args.archive_format == args.ArchiveFormat.ZIP:
archive_format = archive_generator.CollectionArchiveGenerator.ZIP
file_extension = ".zip"
elif args.archive_format == args.ArchiveFormat.TAR_GZ:
archive_format = archive_generator.CollectionArchiveGenerator.TAR_GZ
file_extension = ".tar.gz"
else:
raise ValueError("Unknown archive format: %s" % args.archive_format)
# Only use the new-style flow archive generator for the flows that
# have the GetFilesArchiveMappings defined.
if mappings:
a_gen = archive_generator.FlowArchiveGenerator(flow_object,
archive_format)
content_generator = self._WrapContentGeneratorWithMappings(
a_gen, mappings, args, context=context)
else:
a_gen = archive_generator.CollectionArchiveGenerator(
prefix=target_file_prefix,
description=description,
archive_format=archive_format,
predicate=self._BuildPredicate(str(args.client_id), context=context),
client_id=args.client_id.ToString())
content_generator = self._WrapContentGenerator(
a_gen, flow_results, args, context=context)
return api_call_handler_base.ApiBinaryStream(
target_file_prefix + file_extension,
content_generator=content_generator)
