Skip to content

route zip archives by all valid PK signatures #46

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
abdul-khaliq-khalid wants to merge 1 commit into
base: main
Choose a base branch
from
+12 −9
Open

route zip archives by all valid PK signatures #46

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
5 changes: 5 additions & 0 deletions lib/constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,11 @@
# Combination of all the globals for exclusion to identify unknown method calls
ALL_STRINGS = UNSAFE_STRINGS.union(SUSPICIOUS_STRINGS).union(SAFE_STRINGS)

# ZIP local file header, empty-archive end-of-central-directory record, and the
# spanned-archive marker. zipfile opens archives beginning with any of these
# since it reads the central directory at the end of the file.
ZIP_MAGIC_BYTES = (b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08")

NON_PICKLE_MAGIC_BYTES = (
b"\x7fELF", # ELF executable
b"MZ", # PE executable
Expand Down
4 changes: 1 addition & 3 deletions lib/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,9 +293,7 @@ def is_zip_bytes(file_bytes: bytes | BinaryIO) -> bool:
"""
if not file_bytes:
return False
return _peek_bytes(file_bytes, 4).startswith(
(b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08")
)
return _peek_bytes(file_bytes, 4).startswith(constants.ZIP_MAGIC_BYTES)


def extract_zip_contents(
Expand Down
12 changes: 6 additions & 6 deletions saferpickle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -978,14 +978,14 @@ def security_scan(
header = pickle_bytes.read(262)
pickle_bytes.seek(current_pos)
if header.startswith(
(b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b")
constants.ZIP_MAGIC_BYTES + (b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b")
) or (len(header) >= 262 and header[257:262] == b"ustar"):
is_archive = True

if is_archive:
# Temporarily archive streams fully to bytes
archive_bytes = pickle_bytes.read() # pyrefly: ignore[missing-attribute]
if archive_bytes.startswith(b"PK\x03\x04"):
if archive_bytes.startswith(constants.ZIP_MAGIC_BYTES):
archive_type = "zip"
elif archive_bytes.startswith(b"BZh"):
archive_type = "bz2"
Expand All @@ -1006,7 +1006,7 @@ def security_scan(

# Check for compression signatures if input was raw bytes
if isinstance(pickle_bytes, bytes):
if pickle_bytes.startswith(b"PK\x03\x04"):
if pickle_bytes.startswith(constants.ZIP_MAGIC_BYTES):
return _extract_and_scan_archive(
pickle_bytes,
"zip",
Expand Down Expand Up @@ -1250,7 +1250,7 @@ def _security_scan_internal(
header_bytes = stream.read(1024)
stream.seek(current_pos)
is_archive = header_bytes.startswith(
(b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b")
constants.ZIP_MAGIC_BYTES + (b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b")
) or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar")
if not is_archive:
start_offset = utils.find_pickle_start_offset(stream)
Expand Down Expand Up @@ -1414,7 +1414,7 @@ def _scan_and_load(
pickle_file.seek(0)
header_bytes = pickle_file.read(1024)
is_archive = header_bytes.startswith(
(b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b")
constants.ZIP_MAGIC_BYTES + (b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b")
) or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar")
if not is_archive:
start_offset = utils.find_pickle_start_offset(header_bytes)
Expand All @@ -1434,7 +1434,7 @@ def _scan_and_load(
raise TypeError("pickle_file_or_bytes must be bytes when is_load=False")
data_bytes = pickle_file_or_bytes
is_archive = data_bytes.startswith(
(b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b")
constants.ZIP_MAGIC_BYTES + (b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b")
) or (len(data_bytes) >= 262 and data_bytes[257:262] == b"ustar")
if not is_archive:
start_offset = utils.find_pickle_start_offset(data_bytes)
Expand Down