# We can test part of the module without zlib.

try:

    import zlib

except ImportError:

    zlib = None

import zipfile, os, unittest, sys, shutil, struct

from StringIO import StringIO

from tempfile import TemporaryFile

from random import randint, random

import test.test_support as support

from test.test_support import TESTFN, run_unittest

TESTFN2 = TESTFN + "2"

FIXEDTEST_SIZE = 1000

SMALL_TEST_DATA = [('_ziptest1', '1q2w3e4r5t'),

                   ('ziptest2dir/_ziptest2', 'qawsedrftg'),

                   ('/ziptest2dir/ziptest3dir/_ziptest3', 'azsxdcfvgb'),

                   ('ziptest2dir/ziptest3dir/ziptest4dir/_ziptest3', '6y7u8i9o0p')]

class TestsWithSourceFile(unittest.TestCase):

    def setUp(self):

        self.line_gen = ["Zipfile test line %d. random float: %f" % (i, random())

                          for i in xrange(FIXEDTEST_SIZE)]

        self.data = '\n'.join(self.line_gen) + '\n'

        # Make a source file with some lines

        fp = open(TESTFN, "wb")

        fp.write(self.data)

        fp.close()

    def makeTestArchive(self, f, compression):

        # Create the ZIP archive

        zipfp = zipfile.ZipFile(f, "w", compression)

        zipfp.write(TESTFN, "another"+os.extsep+"name")

        zipfp.write(TESTFN, TESTFN)

        zipfp.writestr("strfile", self.data)

        zipfp.close()

    def zipTest(self, f, compression):

        self.makeTestArchive(f, compression)

        # Read the ZIP archive

        zipfp = zipfile.ZipFile(f, "r", compression)

        self.assertEqual(zipfp.read(TESTFN), self.data)

        self.assertEqual(zipfp.read("another"+os.extsep+"name"), self.data)

        self.assertEqual(zipfp.read("strfile"), self.data)

        # Print the ZIP directory

        fp = StringIO()

        stdout = sys.stdout

        try:

            sys.stdout = fp

            zipfp.printdir()

        finally:

            sys.stdout = stdout

        directory = fp.getvalue()

        lines = directory.splitlines()

        self.assertEquals(len(lines), 4) # Number of files + header

        self.assert_('File Name' in lines[0])

        self.assert_('Modified' in lines[0])

        self.assert_('Size' in lines[0])

        fn, date, time, size = lines[1].split()

        self.assertEquals(fn, 'another.name')

        # XXX: timestamp is not tested

        self.assertEquals(size, str(len(self.data)))

        # Check the namelist

        names = zipfp.namelist()

        self.assertEquals(len(names), 3)

        self.assert_(TESTFN in names)

        self.assert_("another"+os.extsep+"name" in names)

        self.assert_("strfile" in names)

        # Check infolist

        infos = zipfp.infolist()

        names = [ i.filename for i in infos ]

        self.assertEquals(len(names), 3)

        self.assert_(TESTFN in names)

        self.assert_("another"+os.extsep+"name" in names)

        self.assert_("strfile" in names)

        for i in infos:

            self.assertEquals(i.file_size, len(self.data))

        # check getinfo

        for nm in (TESTFN, "another"+os.extsep+"name", "strfile"):

            info = zipfp.getinfo(nm)

            self.assertEquals(info.filename, nm)

            self.assertEquals(info.file_size, len(self.data))

        # Check that testzip doesn't raise an exception

        zipfp.testzip()

        zipfp.close()

    def testStored(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.zipTest(f, zipfile.ZIP_STORED)

    def zipOpenTest(self, f, compression):

        self.makeTestArchive(f, compression)

        # Read the ZIP archive

        zipfp = zipfile.ZipFile(f, "r", compression)

        zipdata1 = []

        zipopen1 = zipfp.open(TESTFN)

        while 1:

            read_data = zipopen1.read(256)

            if not read_data:

                break

            zipdata1.append(read_data)

        zipdata2 = []

        zipopen2 = zipfp.open("another"+os.extsep+"name")

        while 1:

            read_data = zipopen2.read(256)

            if not read_data:

                break

            zipdata2.append(read_data)

        self.assertEqual(''.join(zipdata1), self.data)

        self.assertEqual(''.join(zipdata2), self.data)

        zipfp.close()

    def testOpenStored(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.zipOpenTest(f, zipfile.ZIP_STORED)

    def testOpenViaZipInfo(self):

        # Create the ZIP archive

        zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED)

        zipfp.writestr("name", "foo")

        zipfp.writestr("name", "bar")

        zipfp.close()

        zipfp = zipfile.ZipFile(TESTFN2, "r")

        infos = zipfp.infolist()

        data = ""

        for info in infos:

            data += zipfp.open(info).read()

        self.assert_(data == "foobar" or data == "barfoo")

        data = ""

        for info in infos:

            data += zipfp.read(info)

        self.assert_(data == "foobar" or data == "barfoo")

        zipfp.close()

    def zipRandomOpenTest(self, f, compression):

        self.makeTestArchive(f, compression)

        # Read the ZIP archive

        zipfp = zipfile.ZipFile(f, "r", compression)

        zipdata1 = []

        zipopen1 = zipfp.open(TESTFN)

        while 1:

            read_data = zipopen1.read(randint(1, 1024))

            if not read_data:

                break

            zipdata1.append(read_data)

        self.assertEqual(''.join(zipdata1), self.data)

        zipfp.close()

    def testRandomOpenStored(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.zipRandomOpenTest(f, zipfile.ZIP_STORED)

    def zipReadlineTest(self, f, compression):

        self.makeTestArchive(f, compression)

        # Read the ZIP archive

        zipfp = zipfile.ZipFile(f, "r")

        zipopen = zipfp.open(TESTFN)

        for line in self.line_gen:

            linedata = zipopen.readline()

            self.assertEqual(linedata, line + '\n')

        zipfp.close()

    def zipReadlinesTest(self, f, compression):

        self.makeTestArchive(f, compression)

        # Read the ZIP archive

        zipfp = zipfile.ZipFile(f, "r")

        ziplines = zipfp.open(TESTFN).readlines()

        for line, zipline in zip(self.line_gen, ziplines):

            self.assertEqual(zipline, line + '\n')

        zipfp.close()

    def zipIterlinesTest(self, f, compression):

        self.makeTestArchive(f, compression)

        # Read the ZIP archive

        zipfp = zipfile.ZipFile(f, "r")

        for line, zipline in zip(self.line_gen, zipfp.open(TESTFN)):

            self.assertEqual(zipline, line + '\n')

        zipfp.close()

    def testReadlineStored(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.zipReadlineTest(f, zipfile.ZIP_STORED)

    def testReadlinesStored(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.zipReadlinesTest(f, zipfile.ZIP_STORED)

    def testIterlinesStored(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.zipIterlinesTest(f, zipfile.ZIP_STORED)

    if zlib:

        def testDeflated(self):

            for f in (TESTFN2, TemporaryFile(), StringIO()):

                self.zipTest(f, zipfile.ZIP_DEFLATED)

        def testOpenDeflated(self):

            for f in (TESTFN2, TemporaryFile(), StringIO()):

                self.zipOpenTest(f, zipfile.ZIP_DEFLATED)

        def testRandomOpenDeflated(self):

            for f in (TESTFN2, TemporaryFile(), StringIO()):

                self.zipRandomOpenTest(f, zipfile.ZIP_DEFLATED)

        def testReadlineDeflated(self):

            for f in (TESTFN2, TemporaryFile(), StringIO()):

                self.zipReadlineTest(f, zipfile.ZIP_DEFLATED)

        def testReadlinesDeflated(self):

            for f in (TESTFN2, TemporaryFile(), StringIO()):

                self.zipReadlinesTest(f, zipfile.ZIP_DEFLATED)

        def testIterlinesDeflated(self):

            for f in (TESTFN2, TemporaryFile(), StringIO()):

                self.zipIterlinesTest(f, zipfile.ZIP_DEFLATED)

        def testLowCompression(self):

            # Checks for cases where compressed data is larger than original

            # Create the ZIP archive

            zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_DEFLATED)

            zipfp.writestr("strfile", '12')

            zipfp.close()

            # Get an open object for strfile

            zipfp = zipfile.ZipFile(TESTFN2, "r", zipfile.ZIP_DEFLATED)

            openobj = zipfp.open("strfile")

            self.assertEqual(openobj.read(1), '1')

            self.assertEqual(openobj.read(1), '2')

    def testAbsoluteArcnames(self):

        zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED)

        zipfp.write(TESTFN, "/absolute")

        zipfp.close()

        zipfp = zipfile.ZipFile(TESTFN2, "r", zipfile.ZIP_STORED)

        self.assertEqual(zipfp.namelist(), ["absolute"])

        zipfp.close()

    def testAppendToZipFile(self):

        # Test appending to an existing zipfile

        zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED)

        zipfp.write(TESTFN, TESTFN)

        zipfp.close()

        zipfp = zipfile.ZipFile(TESTFN2, "a", zipfile.ZIP_STORED)

        zipfp.writestr("strfile", self.data)

        self.assertEqual(zipfp.namelist(), [TESTFN, "strfile"])

        zipfp.close()

    def testAppendToNonZipFile(self):

        # Test appending to an existing file that is not a zipfile

        # NOTE: this test fails if len(d) < 22 because of the first

        # line "fpin.seek(-22, 2)" in _EndRecData

        d = 'I am not a ZipFile!'*10

        f = file(TESTFN2, 'wb')

        f.write(d)

        f.close()

        zipfp = zipfile.ZipFile(TESTFN2, "a", zipfile.ZIP_STORED)

        zipfp.write(TESTFN, TESTFN)

        zipfp.close()

        f = file(TESTFN2, 'rb')

        f.seek(len(d))

        zipfp = zipfile.ZipFile(f, "r")

        self.assertEqual(zipfp.namelist(), [TESTFN])

        zipfp.close()

        f.close()

    def test_WriteDefaultName(self):

        # Check that calling ZipFile.write without arcname specified produces the expected result

        zipfp = zipfile.ZipFile(TESTFN2, "w")

        zipfp.write(TESTFN)

        self.assertEqual(zipfp.read(TESTFN), file(TESTFN).read())

        zipfp.close()

    def test_PerFileCompression(self):

        # Check that files within a Zip archive can have different compression options

        zipfp = zipfile.ZipFile(TESTFN2, "w")

        zipfp.write(TESTFN, 'storeme', zipfile.ZIP_STORED)

        zipfp.write(TESTFN, 'deflateme', zipfile.ZIP_DEFLATED)

        sinfo = zipfp.getinfo('storeme')

        dinfo = zipfp.getinfo('deflateme')

        self.assertEqual(sinfo.compress_type, zipfile.ZIP_STORED)

        self.assertEqual(dinfo.compress_type, zipfile.ZIP_DEFLATED)

        zipfp.close()

    def test_WriteToReadonly(self):

        # Check that trying to call write() on a readonly ZipFile object

        # raises a RuntimeError

        zipf = zipfile.ZipFile(TESTFN2, mode="w")

        zipf.writestr("somefile.txt", "bogus")

        zipf.close()

        zipf = zipfile.ZipFile(TESTFN2, mode="r")

        self.assertRaises(RuntimeError, zipf.write, TESTFN)

        zipf.close()

    def testExtract(self):

        zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED)

        for fpath, fdata in SMALL_TEST_DATA:

            zipfp.writestr(fpath, fdata)

        zipfp.close()

        zipfp = zipfile.ZipFile(TESTFN2, "r")

        for fpath, fdata in SMALL_TEST_DATA:

            writtenfile = zipfp.extract(fpath)

            # make sure it was written to the right place

            if os.path.isabs(fpath):

                correctfile = os.path.join(os.getcwd(), fpath[1:])

            else:

                correctfile = os.path.join(os.getcwd(), fpath)

            correctfile = os.path.normpath(correctfile)

            self.assertEqual(writtenfile, correctfile)

            # make sure correct data is in correct file

            self.assertEqual(fdata, file(writtenfile, "rb").read())

            os.remove(writtenfile)

        zipfp.close()

        # remove the test file subdirectories

        shutil.rmtree(os.path.join(os.getcwd(), 'ziptest2dir'))

    def testExtractAll(self):

        zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED)

        for fpath, fdata in SMALL_TEST_DATA:

            zipfp.writestr(fpath, fdata)

        zipfp.close()

        zipfp = zipfile.ZipFile(TESTFN2, "r")

        zipfp.extractall()

        for fpath, fdata in SMALL_TEST_DATA:

            if os.path.isabs(fpath):

                outfile = os.path.join(os.getcwd(), fpath[1:])

            else:

                outfile = os.path.join(os.getcwd(), fpath)

            self.assertEqual(fdata, file(outfile, "rb").read())

            os.remove(outfile)

        zipfp.close()

        # remove the test file subdirectories

        shutil.rmtree(os.path.join(os.getcwd(), 'ziptest2dir'))

    def zip_test_writestr_permissions(self, f, compression):

        # Make sure that writestr creates files with mode 0600,

        # when it is passed a name rather than a ZipInfo instance.

        self.makeTestArchive(f, compression)

        zipfp = zipfile.ZipFile(f, "r")

        zinfo = zipfp.getinfo('strfile')

        self.assertEqual(zinfo.external_attr, 0600 << 16)

    def test_writestr_permissions(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.zip_test_writestr_permissions(f, zipfile.ZIP_STORED)

    def tearDown(self):

        os.remove(TESTFN)

        os.remove(TESTFN2)

class TestZip64InSmallFiles(unittest.TestCase):

    # These tests test the ZIP64 functionality without using large files,

    # see test_zipfile64 for proper tests.

    def setUp(self):

        self._limit = zipfile.ZIP64_LIMIT

        zipfile.ZIP64_LIMIT = 5

        line_gen = ("Test of zipfile line %d." % i for i in range(0, FIXEDTEST_SIZE))

        self.data = '\n'.join(line_gen)

        # Make a source file with some lines

        fp = open(TESTFN, "wb")

        fp.write(self.data)

        fp.close()

    def largeFileExceptionTest(self, f, compression):

        zipfp = zipfile.ZipFile(f, "w", compression)

        self.assertRaises(zipfile.LargeZipFile,

                zipfp.write, TESTFN, "another"+os.extsep+"name")

        zipfp.close()

    def largeFileExceptionTest2(self, f, compression):

        zipfp = zipfile.ZipFile(f, "w", compression)

        self.assertRaises(zipfile.LargeZipFile,

                zipfp.writestr, "another"+os.extsep+"name", self.data)

        zipfp.close()

    def testLargeFileException(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.largeFileExceptionTest(f, zipfile.ZIP_STORED)

            self.largeFileExceptionTest2(f, zipfile.ZIP_STORED)

    def zipTest(self, f, compression):

        # Create the ZIP archive

        zipfp = zipfile.ZipFile(f, "w", compression, allowZip64=True)

        zipfp.write(TESTFN, "another"+os.extsep+"name")

        zipfp.write(TESTFN, TESTFN)

        zipfp.writestr("strfile", self.data)

        zipfp.close()

        # Read the ZIP archive

        zipfp = zipfile.ZipFile(f, "r", compression)

        self.assertEqual(zipfp.read(TESTFN), self.data)

        self.assertEqual(zipfp.read("another"+os.extsep+"name"), self.data)

        self.assertEqual(zipfp.read("strfile"), self.data)

        # Print the ZIP directory

        fp = StringIO()

        stdout = sys.stdout

        try:

            sys.stdout = fp

            zipfp.printdir()

        finally:

            sys.stdout = stdout

        directory = fp.getvalue()

        lines = directory.splitlines()

        self.assertEquals(len(lines), 4) # Number of files + header

        self.assert_('File Name' in lines[0])

        self.assert_('Modified' in lines[0])

        self.assert_('Size' in lines[0])

        fn, date, time, size = lines[1].split()

        self.assertEquals(fn, 'another.name')

        # XXX: timestamp is not tested

        self.assertEquals(size, str(len(self.data)))

        # Check the namelist

        names = zipfp.namelist()

        self.assertEquals(len(names), 3)

        self.assert_(TESTFN in names)

        self.assert_("another"+os.extsep+"name" in names)

        self.assert_("strfile" in names)

        # Check infolist

        infos = zipfp.infolist()

        names = [ i.filename for i in infos ]

        self.assertEquals(len(names), 3)

        self.assert_(TESTFN in names)

        self.assert_("another"+os.extsep+"name" in names)

        self.assert_("strfile" in names)

        for i in infos:

            self.assertEquals(i.file_size, len(self.data))

        # check getinfo

        for nm in (TESTFN, "another"+os.extsep+"name", "strfile"):

            info = zipfp.getinfo(nm)

            self.assertEquals(info.filename, nm)

            self.assertEquals(info.file_size, len(self.data))

        # Check that testzip doesn't raise an exception

        zipfp.testzip()

        zipfp.close()

    def testStored(self):

        for f in (TESTFN2, TemporaryFile(), StringIO()):

            self.zipTest(f, zipfile.ZIP_STORED)

    if zlib:

        def testDeflated(self):

            for f in (TESTFN2, TemporaryFile(), StringIO()):

                self.zipTest(f, zipfile.ZIP_DEFLATED)

    def testAbsoluteArcnames(self):

        zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED, allowZip64=True)

        zipfp.write(TESTFN, "/absolute")

        zipfp.close()

        zipfp = zipfile.ZipFile(TESTFN2, "r", zipfile.ZIP_STORED)

        self.assertEqual(zipfp.namelist(), ["absolute"])

        zipfp.close()

    def tearDown(self):

        zipfile.ZIP64_LIMIT = self._limit

        os.remove(TESTFN)

        os.remove(TESTFN2)

class PyZipFileTests(unittest.TestCase):

    def testWritePyfile(self):

        zipfp  = zipfile.PyZipFile(TemporaryFile(), "w")

        fn = __file__

        if fn.endswith('.pyc') or fn.endswith('.pyo'):

            fn = fn[:-1]

        zipfp.writepy(fn)

        bn = os.path.basename(fn)

        self.assert_(bn not in zipfp.namelist())

        self.assert_(bn + 'o' in zipfp.namelist() or bn + 'c' in zipfp.namelist())

        zipfp.close()