def testTarFileWithSymlink(self):
outfd = StringIO.StringIO()
infd1 = StringIO.StringIO("this is a test string")
st1 = os.stat_result(
(0644, 0, 0, 0, 0, 0, len(infd1.getvalue()), 0, 0, 0))
infd2 = StringIO.StringIO("this is another test string")
st2 = os.stat_result(
(0644, 0, 0, 0, 0, 0, len(infd2.getvalue()), 0, 0, 0))
# Write the zip into a file like object.
with utils.StreamingTarWriter(outfd, mode="w:gz") as writer:
writer.WriteFromFD(infd1, "test1.txt", st=st1)
writer.WriteFromFD(infd2, "subdir/test2.txt", st=st2)
writer.WriteSymlink("test1.txt", "test1.txt.link")
writer.WriteSymlink("subdir/test2.txt", "test2.txt.link")
with tarfile.open(fileobj=StringIO.StringIO(outfd.getvalue()),
mode="r") as test_fd:
test_fd.extractall(self.temp_dir)
link_path = os.path.join(self.temp_dir, "test1.txt.link")
self.assertTrue(os.path.islink(link_path))
self.assertEqual(os.readlink(link_path), "test1.txt")
link_path = os.path.join(self.temp_dir, "test2.txt.link")
self.assertTrue(os.path.islink(link_path))
self.assertEqual(os.readlink(link_path), "subdir/test2.txt")
def testTarFileWithMultipleFiles(self):
outfd = StringIO.StringIO()
infd1 = StringIO.StringIO("this is a test string")
st1 = os.stat_result(
(0644, 0, 0, 0, 0, 0, len(infd1.getvalue()), 0, 0, 0))
infd2 = StringIO.StringIO("this is another test string")
st2 = os.stat_result(
(0644, 0, 0, 0, 0, 0, len(infd2.getvalue()), 0, 0, 0))
# Write the tar into a file like object.
with utils.StreamingTarWriter(outfd, mode="w:gz") as writer:
writer.WriteFromFD(infd1, "test1.txt", st=st1)
writer.WriteFromFD(infd2, "subdir/test2.txt", st=st2)
test_tar = tarfile.open(fileobj=StringIO.StringIO(outfd.getvalue()),
mode="r")
tinfos = sorted(test_tar.getmembers(), key=lambda tinfo: tinfo.name)
self.assertEqual(len(tinfos), 2)
self.assertEqual(tinfos[0].name, "subdir/test2.txt")
self.assertEqual(tinfos[1].name, "test1.txt")
fd = test_tar.extractfile(tinfos[0])
self.assertEqual(fd.read(1024), infd2.getvalue())
fd = test_tar.extractfile(tinfos[1])
self.assertEqual(fd.read(1024), infd1.getvalue())
def testTarFileWithSymlink(self):
outfd = StringIO.StringIO()
infd1 = StringIO.StringIO("this is a test string")
st1 = os.stat_result(
(0644, 0, 0, 0, 0, 0, len(infd1.getvalue()), 0, 0, 0))
infd2 = StringIO.StringIO("this is another test string")
st2 = os.stat_result(
(0644, 0, 0, 0, 0, 0, len(infd2.getvalue()), 0, 0, 0))
# Write the zip into a file like object.
with utils.StreamingTarWriter(outfd, mode="w:gz") as writer:
writer.WriteFromFD(infd1, "test1.txt", st=st1)
writer.WriteFromFD(infd2, "subdir/test2.txt", st=st2)
writer.WriteSymlink("test1.txt", "test1.txt.link")
writer.WriteSymlink("subdir/test2.txt", "test2.txt.link")
with utils.TempDirectory() as temp_dir:
tar_path = os.path.join(temp_dir, "archive.tar.gz")
with open(tar_path, "w") as fd:
fd.write(outfd.getvalue())
# Builtin python ZipFile implementation doesn't support symlinks,
# so we have to extract the files with command line tool.
subprocess.check_call(["tar", "-xzf", tar_path, "-C", temp_dir])
link_path = os.path.join(temp_dir, "test1.txt.link")
self.assertTrue(os.path.islink(link_path))
self.assertEqual(os.readlink(link_path), "test1.txt")
link_path = os.path.join(temp_dir, "test2.txt.link")
self.assertTrue(os.path.islink(link_path))
self.assertEqual(os.readlink(link_path), "subdir/test2.txt")
def CreateArchive(self, _):
# Create an output zip or tar file in the temp space.
with aff4.FACTORY.Create(None, "TempImageFile", token=self.token) as outfd:
if self.args.format == self.args.ArchiveFormat.ZIP:
file_extension = "zip"
elif self.args.format == self.args.ArchiveFormat.TAR_GZ:
file_extension = "tar.gz"
else:
raise ValueError("Unknown archive format: %s" % self.args.format)
outfd.urn = outfd.urn.Add("%s_%X%X.%s" % (
self.args.target_file_prefix, utils.PRNG.GetULong(),
utils.PRNG.GetULong(), file_extension))
self.Log("Will create output on %s" % outfd.urn)
self.state.output_archive_urn = outfd.urn
collection = aff4.FACTORY.Open(self.args.collection_urn, token=self.token)
buffered_outfd = io.BufferedWriter(RawIOBaseBridge(outfd),
buffer_size=1024 * 1024 * 12)
if self.args.format == self.args.ArchiveFormat.ZIP:
streaming_writer = utils.StreamingZipWriter(buffered_outfd, "w",
zipfile.ZIP_DEFLATED)
elif self.args.format == self.args.ArchiveFormat.TAR_GZ:
streaming_writer = utils.StreamingTarWriter(buffered_outfd, "w:gz")
else:
raise ValueError("Unknown archive format: %s" % self.args.format)
with streaming_writer:
self.DownloadCollectionFiles(collection, streaming_writer,
self.args.target_file_prefix)
self.state.output_size = rdfvalue.ByteSize(outfd.size)
def CreateArchive(self, _):
# Create an output zip or tar file in the temp space.
with aff4.FACTORY.Create(None, "TempImageFile",
token=self.token) as outfd:
friendly_hunt_name = self.args.hunt_urn.Basename().replace(
":", "_")
if self.args.format == self.args.ArchiveFormat.ZIP:
file_extension = "zip"
elif self.args.format == self.args.ArchiveFormat.TAR_GZIP:
file_extension = "tar.gz"
else:
raise ValueError("Unknown archive format: %s" %
self.args.format)
outfd.urn = outfd.urn.Add(
"hunt_%s_%X%X.%s" % (friendly_hunt_name, utils.PRNG.GetULong(),
utils.PRNG.GetULong(), file_extension))
self.Log("Will create output on %s" % outfd.urn)
self.state.output_archive_urn = outfd.urn
hunt = aff4.FACTORY.Open(self.args.hunt_urn,
aff4_type="GRRHunt",
token=self.token)
hunt_output_urn = hunt.state.context.results_collection_urn
collection = aff4.FACTORY.Open(hunt_output_urn,
aff4_type="RDFValueCollection",
token=self.token)
buffered_outfd = io.BufferedWriter(RawIOBaseBridge(outfd),
buffer_size=1024 * 1024 * 12)
if self.args.format == self.args.ArchiveFormat.ZIP:
streaming_writer = utils.StreamingZipWriter(
buffered_outfd, "w", zipfile.ZIP_DEFLATED)
elif self.args.format == self.args.ArchiveFormat.TAR_GZIP:
streaming_writer = utils.StreamingTarWriter(
buffered_outfd, "w:gz")
else:
raise ValueError("Unknown archive format: %s" %
self.args.format)
with streaming_writer:
self.DownloadCollectionFiles(collection, streaming_writer,
friendly_hunt_name)
self.state.output_size = rdfvalue.ByteSize(outfd.size)
def testTarFileWithOneFile(self):
infd = StringIO.StringIO("this is a test string")
st = os.stat_result((0644, 0, 0, 0, 0, 0, len(infd.getvalue()), 0, 0, 0))
# Write the tar into a file like object.
outfd = StringIO.StringIO()
with utils.StreamingTarWriter(outfd, mode="w:gz") as writer:
writer.WriteFromFD(infd, "test.txt", st=st)
test_tar = tarfile.open(
fileobj=StringIO.StringIO(outfd.getvalue()), mode="r")
tinfos = list(test_tar.getmembers())
self.assertEqual(len(tinfos), 1)
self.assertEqual(tinfos[0].name, "test.txt")
fd = test_tar.extractfile(tinfos[0])
self.assertEqual(fd.read(1024), infd.getvalue())
