feat: try to guess file type rather than trust file ext

src/picklescan/scanner.py

@@ -283,6 +283,16 @@ def _is_7z_file(f: IO[bytes]) -> bool:
     return read_bytes == local_header_magic_number
 
 
+def _is_legacy_pytorch(data: IO[bytes]) -> bool:
+    start = data.tell()
+    try:
+        magic = get_magic_number(data)
+    except Exception:
+        magic = None
+    data.seek(start)
+    return magic == MAGIC_NUMBER
+
+
 def _http_get(url) -> bytes:
     _log.debug(f"Request: GET {url}")
 
@@ -322,7 +332,12 @@ def _list_globals(data: IO[bytes], multiple_pickles=True) -> Set[Tuple[str, str]
             _log.debug(f"Error parsing pickle: {e}", exc_info=True)
             parsing_pkl_error = str(e)
         last_byte = data.read(1)
-        data.seek(-1, 1)
+        if last_byte != b"":
+            # Rewind the peeked byte so the next genops pass starts on it. A backward seek on a ZipExtFile resets the member and re-reads
+            # it from the start, and on Python 3.12+ that reset also re-enables the CRC validation RelaxedZipFile deliberately disabled
+            # (zipfile restores _expected_crc from a copy captured in ZipExtFile.__init__, before RelaxedZipFile nulled it).
+            # A corrupt-CRC member would then raise BadZipFile here, after parsing succeeded, discarding the globals already found in it.
+            data.seek(-1, 1)
 
         # Extract global imports
         for n in range(len(ops)):
@@ -458,6 +473,9 @@ def scan_pickle_bytes(data: IO[bytes], file_id, multiple_pickles=True, strict=Fa
 
     try:
         raw_globals = _list_globals(data, multiple_pickles)
+    except ValueError as e:
+        _log.warning(f"WARNING: could not parse {file_id} as pickle: {e}")
+        return ScanResult([], scanned_files=1, scan_err=True)
     except GenOpsError as e:
         if e.globals is not None:
             # Found some globals before error - could be a malicious partial pickle
@@ -515,12 +533,18 @@ def scan_zip_bytes(data: IO[bytes], file_id, strict=False) -> ScanResult:
                     magic_bytes = file.read(8)
                 file_ext = os.path.splitext(file_name)[1]
 
-                if file_ext in _pickle_file_extensions or any(magic_bytes.startswith(mn) for mn in _pickle_magic_bytes):
+                # Magic bytes take precedence over the file extension
+                if magic_bytes.startswith(_numpy_magic_bytes):
+                    _log.debug("Scanning file %s in zip archive %s", file_name, file_id)
+                    with zip.open(file_name, "r") as file:
+                        result.merge(scan_numpy(file, f"{file_id}:{file_name}", strict=strict))
+
+                elif file_ext in _pickle_file_extensions or any(magic_bytes.startswith(mn) for mn in _pickle_magic_bytes):
                     _log.debug("Scanning file %s in zip archive %s", file_name, file_id)
                     with zip.open(file_name, "r") as file:
                         result.merge(scan_pickle_bytes(file, f"{file_id}:{file_name}", strict=strict))
 
-                elif file_ext in _numpy_file_extensions or magic_bytes.startswith(_numpy_magic_bytes):
+                elif file_ext in _numpy_file_extensions:
                     _log.debug("Scanning file %s in zip archive %s", file_name, file_id)
                     with zip.open(file_name, "r") as file:
                         result.merge(scan_numpy(file, f"{file_id}:{file_name}", strict=strict))
@@ -546,8 +570,7 @@ def scan_numpy(data: IO[bytes], file_id, strict=False) -> ScanResult:
     # to seek past the beginning of the file
     data.seek(-min(N, len(magic)), 1)  # back-up
     if magic.startswith(_ZIP_PREFIX) or magic.startswith(_ZIP_SUFFIX):
-        # .npz file
-        raise ValueError(f".npz file not handled as zip file: {file_id}")
+        return scan_zip_bytes(data, file_id, strict=strict)
     elif magic == np.lib.format.MAGIC_PREFIX:
         # .npy file
 
@@ -616,6 +639,13 @@ def scan_pytorch(data: IO[bytes], file_id, strict=False) -> ScanResult:
 def scan_bytes(data: IO[bytes], file_id, file_ext: Optional[str] = None, strict=False) -> ScanResult:
     _log.debug(f"scan_bytes({file_id})")
 
+    start = data.tell()
+    magic_bytes = data.read(8)
+    data.seek(start)
+
+    if magic_bytes.startswith(_numpy_magic_bytes):
+        return scan_numpy(data, file_id, strict=strict)
+
     if file_ext is not None and file_ext in _pytorch_file_extensions:
         try:
             return scan_pytorch(data, file_id, strict=strict)
@@ -624,17 +654,21 @@ def scan_bytes(data: IO[bytes], file_id, file_ext: Optional[str] = None, strict=
                 f"WARNING: Invalid PyTorch magic number for file {e}. Trying to scan as non-PyTorch file.",
                 exc_info=_log.isEnabledFor(logging.DEBUG),
             )
-            data.seek(0)
+            data.seek(start)
 
-    if file_ext is not None and file_ext in _numpy_file_extensions:
-        return scan_numpy(data, file_id, strict=strict)
+    if any(magic_bytes.startswith(mn) for mn in _pickle_magic_bytes):
+        if _is_legacy_pytorch(data):
+            return scan_pytorch(data, file_id, strict=strict)
+        return scan_pickle_bytes(data, file_id, strict=strict)
 
     is_zip = zipfile.is_zipfile(data)
-    data.seek(0)
+    data.seek(start)
     if is_zip:
         return scan_zip_bytes(data, file_id, strict=strict)
     elif _is_7z_file(data):
         return scan_7z_bytes(data, file_id, strict=strict)
+    elif _is_legacy_pytorch(data):
+        return scan_pytorch(data, file_id, strict=strict)
     else:
         return scan_pickle_bytes(data, file_id, strict=strict)
 

tests/test_scanner.py

@@ -20,6 +20,7 @@
     _http_get,
     _list_globals,
     _build_scan_result_from_raw_globals,
+    scan_bytes,
     scan_pickle_bytes,
     scan_zip_bytes,
     scan_directory_path,
@@ -782,6 +783,100 @@ def test_scan_file_path_npz():
     )
 
 
+def test_scan_bytes_content_over_extension():
+    """Files whose extension does not match their content are dispatched on magic bytes, not extension (common with mislabeled uploads)."""
+    npy_result = ScanResult(
+        [
+            Global("numpy.core.multiarray", "_reconstruct", SafetyLevel.Innocuous),
+            Global("numpy", "ndarray", SafetyLevel.Innocuous),
+            Global("numpy", "dtype", SafetyLevel.Innocuous),
+        ],
+        scanned_files=1,
+        issues_count=0,
+        infected_files=0,
+    )
+    with open(f"{_root_path}/data2/object_array.npy", "rb") as f:
+        npy_bytes = f.read()
+    # numpy content wearing a pickle/pytorch/no extension must reach scan_numpy, not crash the pickle parser on \x93 (STACK_GLOBAL)
+    for file_ext in (".pkl", ".bin", None):
+        compare_scan_results(scan_bytes(io.BytesIO(npy_bytes), f"mislabeled{file_ext}", file_ext), npy_result)
+
+    # npz (zip) content wearing a .npy extension is scanned as a zip archive
+    with open(f"{_root_path}/data2/object_arrays.npz", "rb") as f:
+        npz_bytes = f.read()
+    npz_result = ScanResult(
+        [
+            Global("numpy.core.multiarray", "_reconstruct", SafetyLevel.Innocuous),
+            Global("numpy", "ndarray", SafetyLevel.Innocuous),
+            Global("numpy", "dtype", SafetyLevel.Innocuous),
+        ]
+        * 2,
+        scanned_files=2,
+        issues_count=0,
+        infected_files=0,
+    )
+    compare_scan_results(scan_bytes(io.BytesIO(npz_bytes), "mislabeled.npy", ".npy"), npz_result)
+    # same content handed straight to scan_numpy (e.g. a zip member) routes to the zip scanner instead of raising
+    compare_scan_results(scan_numpy(io.BytesIO(npz_bytes), "mislabeled.npy"), npz_result)
+
+    # pickle content wearing a .npy extension is scanned as pickle
+    compare_scan_results(
+        scan_bytes(io.BytesIO(pickle.dumps(Malicious1())), "mislabeled.npy", ".npy"),
+        ScanResult([Global("builtins", "eval", SafetyLevel.Dangerous)], 1, 1, 1),
+    )
+
+
+def test_scan_zip_bytes_member_content_over_extension():
+    """A zip member with numpy content and a pickle extension is scanned as numpy."""
+    with open(f"{_root_path}/data2/object_array.npy", "rb") as f:
+        npy_bytes = f.read()
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, "w") as zip:
+        zip.writestr("data.pkl", npy_bytes)
+
+    compare_scan_results(
+        scan_zip_bytes(io.BytesIO(buffer.getbuffer()), "test.zip"),
+        ScanResult(
+            [
+                Global("numpy.core.multiarray", "_reconstruct", SafetyLevel.Innocuous),
+                Global("numpy", "ndarray", SafetyLevel.Innocuous),
+                Global("numpy", "dtype", SafetyLevel.Innocuous),
+            ],
+            scanned_files=1,
+            issues_count=0,
+            infected_files=0,
+        ),
+    )
+
+
+def test_scan_file_path_extensionless_pytorch(tmp_path):
+    """PyTorch files without an extension are detected by content: zip format via magic bytes, legacy format via its pickled
+    serialization magic number, so neither ends up on the generic multi-pickle path with scan_err set."""
+    expected = ScanResult(
+        [
+            Global("torch", "FloatStorage", SafetyLevel.Innocuous),
+            Global("collections", "OrderedDict", SafetyLevel.Innocuous),
+            Global("torch._utils", "_rebuild_tensor_v2", SafetyLevel.Innocuous),
+        ],
+        scanned_files=1,
+        issues_count=0,
+        infected_files=0,
+    )
+    for fixture in ("pytorch_model.bin", "new_pytorch_model.bin"):
+        with open(f"{_root_path}/data/{fixture}", "rb") as f:
+            target = tmp_path / f"extensionless_{fixture.split('.')[0]}"
+            target.write_bytes(f.read())
+        result = scan_file_path(str(target))
+        compare_scan_results(result, expected)
+        assert result.scan_err is False
+
+
+def test_stack_global_underflow_does_not_crash():
+    result = scan_pickle_bytes(io.BytesIO(b"\x93NUMPY\x01\x00v\x00"), "not_a_pickle")
+    compare_scan_results(result, ScanResult([], scanned_files=1, issues_count=0, infected_files=0))
+    assert result.scan_err is True
+
+
 def test_scan_directory_path():
     sr = ScanResult(
         globals=[