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())