def testGeneratesTarGzArchive(self):
result = self.handler.Handle(hunt_plugin.ApiGetHuntFilesArchiveArgs(
hunt_id=self.hunt.urn.Basename(), archive_format="TAR_GZ"),
token=self.token)
with utils.TempDirectory() as temp_dir:
tar_path = os.path.join(temp_dir, "archive.tar.gz")
with open(tar_path, "wb") as fd:
for chunk in result.GenerateContent():
fd.write(chunk)
with tarfile.open(tar_path) as tar_fd:
tar_fd.extractall(path=temp_dir)
manifest_file_path = None
for parent, _, files in os.walk(temp_dir):
if "MANIFEST" in files:
manifest_file_path = os.path.join(parent, "MANIFEST")
break
self.assertTrue(manifest_file_path)
with open(manifest_file_path, "rb") as fd:
manifest = yaml.safe_load(fd.read())
self.assertEqual(manifest["archived_files"], 10)
self.assertEqual(manifest["failed_files"], 0)
self.assertEqual(manifest["processed_files"], 10)
self.assertEqual(manifest["ignored_files"], 0)
def testGeneratesZipArchive(self):
result = self.handler.Handle(hunt_plugin.ApiGetHuntFilesArchiveArgs(
hunt_id=self.hunt.urn.Basename(), archive_format="ZIP"),
token=self.token)
out_fd = StringIO.StringIO()
for chunk in result.GenerateContent():
out_fd.write(chunk)
zip_fd = zipfile.ZipFile(out_fd, "r")
manifest = None
for name in zip_fd.namelist():
if name.endswith("MANIFEST"):
manifest = yaml.safe_load(zip_fd.read(name))
self.assertEqual(manifest["archived_files"], 10)
self.assertEqual(manifest["failed_files"], 0)
self.assertEqual(manifest["processed_files"], 10)
self.assertEqual(manifest["ignored_files"], 0)
def testGetHuntFilesArchiveIsAccessChecked(self):
args = api_hunt.ApiGetHuntFilesArchiveArgs(hunt_id="H:123456")
self.CheckMethodIsAccessChecked(self.router.GetHuntFilesArchive,
"CheckHuntAccess",
args=args)