Skip to content

Scan compressed pickle artifacts #68

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
massy-o wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Scan compressed pickle artifacts #68

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 42 additions & 11 deletions src/picklescan/scanner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
import bz2
from dataclasses import dataclass, field
from enum import Enum
import gzip
import http.client
import io
import json
import lzma
import logging
import os
import pickletools
Expand Down Expand Up @@ -246,6 +249,17 @@ def __str__(self) -> str:
_pytorch_file_extensions = {".bin", ".pt", ".pth", ".ckpt"}
_pickle_file_extensions = {".pkl", ".pickle", ".joblib", ".dat", ".data"}
_zip_file_extensions = {".zip", ".npz", ".7z"}
_compressed_pickle_suffixes = {
".bz2": bz2.decompress,
".gz": gzip.decompress,
".lzma": lzma.decompress,
".xz": lzma.decompress,
}
_compressed_pickle_file_extensions = {

@massy-o massy-o May 14, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review: compound extensions are generated from the existing pickle extension allowlist, so this expands coverage for compressed forms without making every compressed file type scan as pickle.

f"{pickle_ext}{compressed_ext}" for pickle_ext in _pickle_file_extensions for compressed_ext in _compressed_pickle_suffixes
}
_all_pickle_file_extensions = _pickle_file_extensions | _compressed_pickle_file_extensions
_supported_file_extensions = _all_pickle_file_extensions | _zip_file_extensions | _pytorch_file_extensions | _numpy_file_extensions
# Pickle files do not actually have magic bytes, but v2+ files
# start with a PROTO (\x80) opcode followed by a byte with the protocol version
_pickle_magic_bytes = {
Expand All @@ -259,6 +273,14 @@ def __str__(self) -> str:
_numpy_magic_bytes = b"\x93NUMPY"


def _get_file_extension(path: str) -> str:
path = str(path)
for file_ext in sorted(_supported_file_extensions, key=len, reverse=True):
if path.endswith(file_ext):
return file_ext
return os.path.splitext(path)[1]


def _is_7z_file(f: IO[bytes]) -> bool:
read_bytes = []
start = f.tell()
Expand Down Expand Up @@ -431,10 +453,19 @@ def _build_scan_result_from_raw_globals(
return ScanResult(globals, 1, issues_count, 1 if issues_count > 0 else 0, scan_err)


def scan_pickle_bytes(data: IO[bytes], file_id, multiple_pickles=True) -> ScanResult:
def scan_pickle_bytes(data: IO[bytes], file_id, multiple_pickles=True, file_ext: Optional[str] = None) -> ScanResult:
"""Disassemble a Pickle stream and report issues"""
_log.debug(f"scan_pickle_bytes({file_id})")

if file_ext in _compressed_pickle_file_extensions:

@massy-o massy-o May 14, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review: decompression happens before pickle opcode parsing and decompression failures are surfaced as scan errors, avoiding the previous clean pass on gzip headers.

compressed_ext = os.path.splitext(file_ext)[1]
decompressor = _compressed_pickle_suffixes[compressed_ext]
try:
data = io.BytesIO(decompressor(data.read()))
except Exception as e:
_log.warning("WARNING: could not decompress %s as %s: %s", file_id, compressed_ext, e)
return ScanResult([], scanned_files=1, scan_err=True)

try:
raw_globals = _list_globals(data, multiple_pickles)
except GenOpsError as e:
Expand Down Expand Up @@ -464,7 +495,7 @@ def scan_7z_bytes(data: IO[bytes], file_id) -> ScanResult:

with py7zr.SevenZipFile(data, mode="r") as archive:
file_names = archive.getnames()
targets = [f for f in file_names if f.endswith(tuple(_pickle_file_extensions))]
targets = [f for f in file_names if _get_file_extension(f) in _all_pickle_file_extensions]
_log.debug("Files in 7z archive %s: %s", file_id, targets)
with TemporaryDirectory() as tmpdir:
archive.extract(path=tmpdir, targets=targets)
Expand All @@ -489,12 +520,12 @@ def scan_zip_bytes(data: IO[bytes], file_id) -> ScanResult:
try:
with zip.open(file_name, "r") as file:
magic_bytes = file.read(8)
file_ext = os.path.splitext(file_name)[1]
file_ext = _get_file_extension(file_name)

if file_ext in _pickle_file_extensions or any(magic_bytes.startswith(mn) for mn in _pickle_magic_bytes):
if file_ext in _all_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}"))
result.merge(scan_pickle_bytes(file, f"{file_id}:{file_name}", file_ext=file_ext))

@massy-o massy-o May 14, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review: this keeps archive member scanning on the same compound-extension path as direct files, which closes the equivalent bypass for .joblib.gz stored inside a zip.


elif file_ext in _numpy_file_extensions or magic_bytes.startswith(_numpy_magic_bytes):
_log.debug("Scanning file %s in zip archive %s", file_name, file_id)
Expand Down Expand Up @@ -612,7 +643,7 @@ def scan_bytes(data: IO[bytes], file_id, file_ext: Optional[str] = None) -> Scan
elif _is_7z_file(data):
return scan_7z_bytes(data, file_id)
else:
return scan_pickle_bytes(data, file_id)
return scan_pickle_bytes(data, file_id, file_ext=file_ext)


def scan_huggingface_model(repo_id):
Expand All @@ -625,8 +656,8 @@ def scan_huggingface_model(repo_id):
# Scan model files
scan_result = ScanResult([])
for file_name in file_names:
file_ext = os.path.splitext(file_name)[1]
if file_ext not in _zip_file_extensions and file_ext not in _pickle_file_extensions and file_ext not in _pytorch_file_extensions:
file_ext = _get_file_extension(file_name)
if file_ext not in _zip_file_extensions and file_ext not in _all_pickle_file_extensions and file_ext not in _pytorch_file_extensions:
continue
_log.debug("Scanning file %s in model %s", file_name, repo_id)
url = f"https://huggingface.co/{repo_id}/resolve/main/{file_name}"
Expand Down Expand Up @@ -662,10 +693,10 @@ def scan_directory_path(path, scan_filter: Optional[ScanFilter] = None) -> ScanR
dir_names[:] = filtered_dirs

for file_name in file_names:
file_ext = os.path.splitext(file_name)[1]
file_ext = _get_file_extension(file_name)
if (
file_ext not in _zip_file_extensions
and file_ext not in _pickle_file_extensions
and file_ext not in _all_pickle_file_extensions
and file_ext not in _pytorch_file_extensions
):
continue
Expand All @@ -690,7 +721,7 @@ def scan_directory_path(path, scan_filter: Optional[ScanFilter] = None) -> ScanR
def scan_file_path(path) -> ScanResult:
_log.debug(f"scan_file_path({path})")

file_ext = os.path.splitext(path)[1]
file_ext = _get_file_extension(path)
with open(path, "rb") as file:
return scan_bytes(file, path, file_ext)

Expand Down
32 changes: 32 additions & 0 deletions tests/test_scanner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import gzip
import http.client
import importlib
import io
Expand Down Expand Up @@ -152,6 +153,37 @@ def test_scan_zip_bytes():
)


def test_scan_compressed_joblib_file_path(tmp_path):

@massy-o massy-o May 14, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review: these regression tests exercise the direct path, directory traversal, and zip-member path so the fix is covered where the scanner previously relied on single-suffix extension checks.

file_path = tmp_path / "model.joblib.gz"
file_path.write_bytes(gzip.compress(pickle.dumps(Malicious2(), protocol=4)))

compare_scan_results(
scan_file_path(str(file_path)),
ScanResult([Global(os.name, "system", SafetyLevel.Dangerous)], 1, 1, 1),
)


def test_scan_directory_path_includes_compressed_joblib(tmp_path):
file_path = tmp_path / "model.joblib.gz"
file_path.write_bytes(gzip.compress(pickle.dumps(Malicious2(), protocol=4)))

compare_scan_results(
scan_directory_path(str(tmp_path)),
ScanResult([Global(os.name, "system", SafetyLevel.Dangerous)], 1, 1, 1),
)


def test_scan_zip_bytes_includes_compressed_joblib_member():
buffer = io.BytesIO()
with zipfile.ZipFile(buffer, "w") as zip:
zip.writestr("model.joblib.gz", gzip.compress(pickle.dumps(Malicious2(), protocol=4)))

compare_scan_results(
scan_zip_bytes(io.BytesIO(buffer.getbuffer()), "test.zip"),
ScanResult([Global(os.name, "system", SafetyLevel.Dangerous)], 1, 1, 1),
)


def test_scan_numpy():
with open(f"{_root_path}/data2/object_array.npy", "rb") as f:
compare_scan_results(
Expand Down