From 0be506a5baf8b86c519e6074af43a0f91af603c6 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Thu, 31 Jan 2013 15:26:55 +0200
Subject: [PATCH] Issue #4844: ZipFile now raises BadZipfile when opens a ZIP
 file with an incomplete "End of Central Directory" record.  Original patch by
 Guilherme Polo and Alan McIntyre.

---
 Lib/test/test_zipfile.py | 14 ++++++++++++++
 Lib/zipfile.py           | 27 ++++++++++++++++++++-------
 Misc/NEWS                |  4 ++++
 3 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py
index 9bd3b5834c2..a5ff90a7346 100644
--- a/Lib/test/test_zipfile.py
+++ b/Lib/test/test_zipfile.py
@@ -760,6 +760,20 @@ class OtherTests(unittest.TestCase):
         chk = zipfile.is_zipfile(fp)
         self.assertTrue(not chk)
 
+    def test_damaged_zipfile(self):
+        """Check that zipfiles with missing bytes at the end raise BadZipFile."""
+        # - Create a valid zip file
+        fp = io.BytesIO()
+        with zipfile.ZipFile(fp, mode="w") as zipf:
+            zipf.writestr("foo.txt", b"O, for a Muse of Fire!")
+        zipfiledata = fp.getvalue()
+
+        # - Now create copies of it missing the last N bytes and make sure
+        #   a BadZipFile exception is raised when we try to open it
+        for N in range(len(zipfiledata)):
+            fp = io.BytesIO(zipfiledata[:N])
+            self.assertRaises(zipfile.BadZipfile, zipfile.ZipFile, fp)
+
     def test_is_zip_valid_file(self):
         """Check that is_zipfile() correctly identifies zip files."""
         # - passing a filename
diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index f340aa65758..9dc4f57489f 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -166,6 +166,8 @@ def _EndRecData64(fpin, offset, endrec):
         return endrec
 
     data = fpin.read(sizeEndCentDir64Locator)
+    if len(data) != sizeEndCentDir64Locator:
+        return endrec
     sig, diskno, reloff, disks = struct.unpack(structEndArchive64Locator, data)
     if sig != stringEndArchive64Locator:
         return endrec
@@ -176,6 +178,8 @@ def _EndRecData64(fpin, offset, endrec):
     # Assume no 'zip64 extensible data'
     fpin.seek(offset - sizeEndCentDir64Locator - sizeEndCentDir64, 2)
     data = fpin.read(sizeEndCentDir64)
+    if len(data) != sizeEndCentDir64:
+        return endrec
     sig, sz, create_version, read_version, disk_num, disk_dir, \
             dircount, dircount2, dirsize, diroffset = \
             struct.unpack(structEndArchive64, data)
@@ -211,7 +215,9 @@ def _EndRecData(fpin):
     except IOError:
         return None
     data = fpin.read()
-    if data[0:4] == stringEndArchive and data[-2:] == "\000\000":
+    if (len(data) == sizeEndCentDir and
+        data[0:4] == stringEndArchive and
+        data[-2:] == b"\000\000"):
         # the signature is correct and there's no comment, unpack structure
         endrec = struct.unpack(structEndArchive, data)
         endrec=list(endrec)
@@ -235,6 +241,9 @@ def _EndRecData(fpin):
     if start >= 0:
         # found the magic number; attempt to unpack and interpret
         recData = data[start:start+sizeEndCentDir]
+        if len(recData) != sizeEndCentDir:
+            # Zip file is corrupted.
+            return None
         endrec = list(struct.unpack(structEndArchive, recData))
         commentSize = endrec[_ECD_COMMENT_SIZE] #as claimed by the zip file
         comment = data[start+sizeEndCentDir:start+sizeEndCentDir+commentSize]
@@ -246,7 +255,7 @@ def _EndRecData(fpin):
                              endrec)
 
     # Unable to find a valid end of central directory structure
-    return
+    return None
 
 
 class ZipInfo (object):
@@ -818,9 +827,11 @@ class ZipFile(object):
         total = 0
         while total < size_cd:
             centdir = fp.read(sizeCentralDir)
-            if centdir[0:4] != stringCentralDir:
-                raise BadZipfile, "Bad magic number for central directory"
+            if len(centdir) != sizeCentralDir:
+                raise BadZipfile("Truncated central directory")
             centdir = struct.unpack(structCentralDir, centdir)
+            if centdir[_CD_SIGNATURE] != stringCentralDir:
+                raise BadZipfile("Bad magic number for central directory")
             if self.debug > 2:
                 print centdir
             filename = fp.read(centdir[_CD_FILENAME_LENGTH])
@@ -948,10 +959,12 @@ class ZipFile(object):
 
             # Skip the file header:
             fheader = zef_file.read(sizeFileHeader)
-            if fheader[0:4] != stringFileHeader:
-                raise BadZipfile, "Bad magic number for file header"
-
+            if len(fheader) != sizeFileHeader:
+                raise BadZipfile("Truncated file header")
             fheader = struct.unpack(structFileHeader, fheader)
+            if fheader[_FH_SIGNATURE] != stringFileHeader:
+                raise BadZipfile("Bad magic number for file header")
+
             fname = zef_file.read(fheader[_FH_FILENAME_LENGTH])
             if fheader[_FH_EXTRA_FIELD_LENGTH]:
                 zef_file.read(fheader[_FH_EXTRA_FIELD_LENGTH])
diff --git a/Misc/NEWS b/Misc/NEWS
index 6a3d9c74c1c..34e4bda9c23 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -202,6 +202,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #4844: ZipFile now raises BadZipfile when opens a ZIP file with an
+  incomplete "End of Central Directory" record.  Original patch by Guilherme
+  Polo and Alan McIntyre.
+
 - Issue #15505: `unittest.installHandler` no longer assumes SIGINT handler is
   set to a callable object.