Improve numpy compatibility and non-pickle handling (#58)

Address #57

Author: mmaitre314

src/picklescan/scanner.py

@@ -400,11 +400,14 @@ def scan_pickle_bytes(data: IO[bytes], file_id, multiple_pickles=True) -> ScanRe
     try:
         raw_globals = _list_globals(data, multiple_pickles)
     except GenOpsError as e:
-        _log.error(f"ERROR: parsing pickle in {file_id}: {e}", exc_info=_log.isEnabledFor(logging.DEBUG))
         if e.globals is not None:
+            # Found some globals before error - could be a malicious partial pickle
+            _log.error(f"ERROR: parsing pickle in {file_id}: {e}", exc_info=_log.isEnabledFor(logging.DEBUG))
             return _build_scan_result_from_raw_globals(e.globals, file_id, scan_err=True)
         else:
-            return ScanResult([], scan_err=True)
+            # No globals found - likely not a pickle file at all
+            _log.warning(f"WARNING: could not parse {file_id} as pickle: {e}")
+            return ScanResult([], scanned_files=1, scan_err=False)
 
     _log.debug("Global imports in %s: %s", file_id, raw_globals)
 
@@ -487,8 +490,12 @@ def scan_numpy(data: IO[bytes], file_id) -> ScanResult:
         # .npy file
 
         version = np.lib.format.read_magic(data)
-        np.lib.format._check_version(version)
-        _, _, dtype = np.lib.format._read_array_header(data, version)
+        if version == (1, 0):
+            _, _, dtype = np.lib.format.read_array_header_1_0(data)
+        elif version in [(2, 0), (3, 0)]:
+            _, _, dtype = np.lib.format.read_array_header_2_0(data)
+        else:
+            raise ValueError(f"Unsupported numpy format version: {version}")
 
         if dtype.hasobject:
             return scan_pickle_bytes(data, file_id)

tests/data/not_a_pickle.bin

@@ -0,0 +1,2 @@
+cmod
+�

tests/init_data_files.py

@@ -469,6 +469,26 @@ def initialize_corrupt_zip_file_crc(path: str, file_name: str, data: bytes):
     print(f"Initialized file {path}.")
 
 
+def initialize_not_a_pickle_file(path: str):
+    """Create a binary file that starts with pickle GLOBAL opcode (0x63 = 'c') but contains
+    invalid UTF-8 bytes, causing 'utf-8' codec can't decode byte error.
+    This reproduces the issue seen with files like vitpose_h_wholebody_data.bin.
+    """
+    if os.path.exists(path):
+        print(f"File {path} already exists, skipping initialization.")
+        return
+
+    # Byte 0x63 is 'c' which is the GLOBAL opcode in pickle protocol 0/1/2
+    # The GLOBAL opcode expects to read two newline-terminated strings (module and name)
+    # We provide a valid module name "mod" followed by newline, then an invalid UTF-8 byte (0xf8)
+    # This triggers: 'utf-8' codec can't decode byte 0xf8 in position 0: invalid start byte
+    data = b"cmod\n\xf8\n"
+
+    with open(path, "wb") as file:
+        file.write(data)
+    print(f"Initialized file {path}.")
+
+
 def initialize_pickle_files():
     os.makedirs(f"{_root_path}/data", exist_ok=True)
 
@@ -796,6 +816,8 @@ def initialize_pickle_files():
     initialize_pickle_file_from_reduce("io_FileIO.pkl", reduce_io_FileIO)
     initialize_pickle_file_from_reduce("urllib_request_urlopen.pkl", reduce_urllib_request_urlopen)
 
+    initialize_not_a_pickle_file(f"{_root_path}/data/not_a_pickle.bin")
+
 
 def initialize_numpy_files():
     import numpy as np

tests/test_scanner.py

@@ -287,7 +287,9 @@ def test_scan_file_path():
     malicious10 = ScanResult([Global("__builtin__", "exec", SafetyLevel.Dangerous)], 1, 1, 1)
     compare_scan_results(scan_file_path(f"{_root_path}/data/malicious10.pkl"), malicious10)
 
-    bad_pytorch = ScanResult([], 0, 0, 0, True)
+    # bad_pytorch.pt is a PNG file with .pt extension - scanner should recognize it's not a valid pickle
+    # and report it as scanned (scanned_files=1) but without errors (scan_err=False) since no threats were found
+    bad_pytorch = ScanResult([], 1, 0, 0, False)
     compare_scan_results(scan_file_path(f"{_root_path}/data/bad_pytorch.pt"), bad_pytorch)
 
     malicious14 = ScanResult([Global("runpy", "_run_code", SafetyLevel.Dangerous)], 1, 1, 1)
@@ -508,9 +510,11 @@ def test_scan_directory_path():
             Global("builtins", "eval", SafetyLevel.Dangerous),
             Global("builtins", "eval", SafetyLevel.Dangerous),
         ],
-        scanned_files=42,
+        scanned_files=44,
         issues_count=43,
         infected_files=37,
+        # scan_err=True because some files (broken_model.pkl, malicious-invalid-bytes.pkl) have partial parsing errors
+        scan_err=True,
     )
     compare_scan_results(scan_directory_path(f"{_root_path}/data/"), sr)
 
@@ -556,3 +560,13 @@ def test_invalid_bytes_err():
             scan_pickle_bytes(file, f"{_root_path}/data/malicious-invalid-bytes.pkl"),
             malicious_invalid_bytes,
         )
+
+
+def test_not_a_pickle_file():
+    """Test scanning a binary file that starts with pickle GLOBAL opcode but has invalid UTF-8.
+    This reproduces the 'utf-8' codec can't decode byte error seen with files like vitpose_h_wholebody_data.bin.
+    The scanner should handle this gracefully: file is scanned, no threats found, no error.
+    """
+    # File is not a valid pickle, but scanner should not error - just report no threats
+    not_a_pickle = ScanResult([], scanned_files=1, issues_count=0, infected_files=0, scan_err=False)
+    compare_scan_results(scan_file_path(f"{_root_path}/data/not_a_pickle.bin"), not_a_pickle)