Merge commit from fork

* more relaxed ZipFile impl

* lint

---------

Co-authored-by: Matthieu Maitre <[email protected]>

Author: Trevor Madge

src/picklescan/relaxed_zipfile.py

@@ -0,0 +1,88 @@
+import struct
+import zipfile
+
+# More forgiving implementation of zipfile.ZipFile
+_FH_SIGNATURE = 0
+_FH_FILENAME_LENGTH = 10
+_FH_EXTRA_FIELD_LENGTH = 11
+
+structFileHeader = "<4s2B4HL2L2H"
+stringFileHeader = b"PK\003\004"
+sizeFileHeader = struct.calcsize(structFileHeader)
+
+
+class RelaxedZipFile(zipfile.ZipFile):
+    def open(self, name, mode="r", pwd=None, *, force_zip64=False):
+        # near copy of zipfile.ZipFile.open with
+        """Return file-like object for 'name'.
+
+        name is a string for the file name within the ZIP file, or a ZipInfo
+        object.
+
+        mode should be 'r' to read a file already in the ZIP file, or 'w' to
+        write to a file newly added to the archive.
+
+        pwd is the password to decrypt files (only used for reading).
+
+        When writing, if the file size is not known in advance but may exceed
+        2 GiB, pass force_zip64 to use the ZIP64 format, which can handle large
+        files.  If the size is known in advance, it is best to pass a ZipInfo
+        instance for name, with zinfo.file_size set.
+        """
+        if mode not in {"r", "w"}:
+            raise ValueError('open() requires mode "r" or "w"')
+        if pwd and not isinstance(pwd, bytes):
+            raise TypeError("pwd: expected bytes, got %s" % type(pwd).__name__)
+        if pwd and (mode == "w"):
+            raise ValueError("pwd is only supported for reading files")
+        if not self.fp:
+            raise ValueError("Attempt to use ZIP archive that was already closed")
+
+        # Make sure we have an info object
+        if isinstance(name, zipfile.ZipInfo):
+            # 'name' is already an info object
+            zinfo = name
+        elif mode == "w":
+            zinfo = zipfile.ZipInfo(name)
+            zinfo.compress_type = self.compression
+            zinfo._compresslevel = self.compresslevel
+        else:
+            # Get info object for name
+            zinfo = self.getinfo(name)
+
+        if mode == "w":
+            return self._open_to_write(zinfo, force_zip64=force_zip64)
+
+        if self._writing:
+            raise ValueError(
+                "Can't read from the ZIP file while there "
+                "is an open writing handle on it. "
+                "Close the writing handle before trying to read."
+            )
+
+        # Open for reading:
+        self._fileRefCnt += 1
+        zef_file = zipfile._SharedFile(
+            self.fp,
+            zinfo.header_offset,
+            self._fpclose,
+            self._lock,
+            lambda: self._writing,
+        )
+        try:
+            # Skip the file header:
+            fheader = zef_file.read(sizeFileHeader)
+            if len(fheader) != sizeFileHeader:
+                raise zipfile.BadZipFile("Truncated file header")
+            fheader = struct.unpack(structFileHeader, fheader)
+            if fheader[_FH_SIGNATURE] != stringFileHeader:
+                raise zipfile.BadZipFile("Bad magic number for file header")
+
+            zef_file.read(fheader[_FH_FILENAME_LENGTH])
+            if fheader[_FH_EXTRA_FIELD_LENGTH]:
+                zef_file.read(fheader[_FH_EXTRA_FIELD_LENGTH])
+
+            return zipfile.ZipExtFile(zef_file, mode, zinfo, pwd, True)
+        except BaseException:
+            zef_file.close()
+            raise

src/picklescan/scanner.py

@@ -11,6 +11,7 @@
 from typing import IO, List, Optional, Set, Tuple
 import urllib.parse
 import zipfile
+from .relaxed_zipfile import RelaxedZipFile
 
 from .torch import (
     get_magic_number,
@@ -375,7 +376,7 @@ def get_magic_bytes_from_zipfile(zip: zipfile.ZipFile, num_bytes=8):
 def scan_zip_bytes(data: IO[bytes], file_id) -> ScanResult:
     result = ScanResult([])
 
-    with zipfile.ZipFile(data, "r") as zip:
+    with RelaxedZipFile(data, "r") as zip:
         magic_bytes = get_magic_bytes_from_zipfile(zip)
         file_names = zip.namelist()
         _log.debug("Files in zip archive %s: %s", file_id, file_names)

tests/data/malicious1_0x1.zip

@@ -249,6 +249,22 @@ def initialize_zip_file(path, file_name, data):
             zip.writestr(file_name, data)
 
 
+def initialize_corrupt_zip_file_central_directory(path, file_name, data):
+    if not os.path.exists(path):
+        with zipfile.ZipFile(path, "w") as zip:
+            zip.writestr(file_name, data)
+
+        with open(path, "rb") as f:
+            data = f.read()
+
+        # Replace only the first occurrence of "data.pkl" with "datap.kl"
+        modified_data = data.replace(b"data.pkl", b"datap.kl", 1)
+
+        # Write back the modified content
+        with open(path, "wb") as f:
+            f.write(modified_data)
+
+
 def initialize_numpy_files():
     import numpy as np
 
@@ -490,6 +506,12 @@ def initialize_pickle_files():
         pickle.dumps(Malicious1(), protocol=4),
     )
 
+    initialize_corrupt_zip_file_central_directory(
+        f"{_root_path}/data/malicious1_central_directory.zip",
+        "data.pkl",
+        pickle.dumps(Malicious1(), protocol=4),
+    )
+
     initialize_zip_file(
         f"{_root_path}/data/malicious1_wrong_ext.zip",
         "data.txt",  # Pickle file with a non-standard extension
@@ -646,7 +668,22 @@ def test_scan_file_path():
     compare_scan_results(
         scan_file_path(f"{_root_path}/data/malicious1.zip"), malicious1
     )
-    compare_scan_results(scan_file_path(f"{_root_path}/data/malicious1.7z"), malicious1)
+    compare_scan_results(
+        scan_file_path(f"{_root_path}/data/malicious1_central_directory.zip"),
+        malicious1,
+    )
+    compare_scan_results(
+        scan_file_path(f"{_root_path}/data/malicious1_0x1.zip"), malicious1
+    )
+    compare_scan_results(
+        scan_file_path(f"{_root_path}/data/malicious1_0x20.zip"), malicious1
+    )
+    compare_scan_results(
+        scan_file_path(f"{_root_path}/data/malicious1_0x40.zip"), malicious1
+    )
+    compare_scan_results(
+        scan_file_path(f"{_root_path}/data/malicious1.7z"), malicious1
+    )
     compare_scan_results(
         scan_file_path(f"{_root_path}/data/malicious1_wrong_ext.zip"), malicious1
     )
@@ -835,6 +872,9 @@ def test_scan_directory_path():
             Global("functools", "partial", SafetyLevel.Dangerous),
             Global("pip", "main", SafetyLevel.Dangerous),
             Global("builtins", "eval", SafetyLevel.Dangerous),
+            Global("builtins", "eval", SafetyLevel.Dangerous),
+            Global("builtins", "eval", SafetyLevel.Dangerous),
+            Global("builtins", "eval", SafetyLevel.Dangerous),
         ],
         scanned_files=38,
         issues_count=39,