diff --git a/lib/utils.py b/lib/utils.py index da97476..6b108f5 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -336,6 +336,42 @@ def is_lzma_bytes(file_bytes: bytes | BinaryIO) -> bool: return _peek_bytes(file_bytes, 6).startswith(b"\xfd\x37\x7a\x58\x5a\x00") +def is_lzma_alone_bytes(file_bytes: bytes | BinaryIO) -> bool: + """Checks if the bytes/stream represent a legacy LZMA (.lzma) alone stream. + + The XZ container carries the b"\\xfd7zXZ\\x00" magic, but the older + FORMAT_ALONE (.lzma) container has no magic bytes. lzma.open() decodes both + formats, so an alone-format stream has to be recognised from the structure of + its 13-byte header or it slips past archive routing. + + Args: + file_bytes: The bytes or stream to check. + + Returns: + True if the header matches a legacy LZMA alone stream, False otherwise. + """ + peeked = _peek_bytes(file_bytes, 13) + if len(peeked) < 13: + return False + # Properties byte encodes (pb * 5 + lp) * 9 + lc; its maximum valid value is + # (4 * 5 + 4) * 9 + 8 == 224. + if peeked[0] > 224: + return False + dict_size = int.from_bytes(peeked[1:5], "little") + # Real encoders always pick a power-of-two dictionary size within the LZMA + # range of 4 KiB to 2 GiB. + if ( + dict_size < (1 << 12) + or dict_size > (1 << 31) + or (dict_size & (dict_size - 1)) != 0 + ): + return False + uncompressed_size = int.from_bytes(peeked[5:13], "little") + # The size field is either the streaming "unknown" sentinel or the real + # length, which comfortably fits in six bytes. + return uncompressed_size == 0xFFFFFFFFFFFFFFFF or uncompressed_size < (1 << 48) + + def extract_lzma_contents(file_bytes: bytes | BinaryIO) -> IO[bytes]: """Extracts contents from lzma bytes/stream as a stream.""" stream = ( diff --git a/saferpickle.py b/saferpickle.py index 974f181..7278a7d 100644 --- a/saferpickle.py +++ b/saferpickle.py @@ -977,9 +977,13 @@ def security_scan( current_pos = pickle_bytes.tell() 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") - ) or (len(header) >= 262 and header[257:262] == b"ustar"): + if ( + header.startswith( + (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") + ) + or (len(header) >= 262 and header[257:262] == b"ustar") + or utils.is_lzma_alone_bytes(header) + ): is_archive = True if is_archive: @@ -993,6 +997,8 @@ def security_scan( archive_type = "lzma" elif archive_bytes.startswith(b"\x1f\x8b"): archive_type = "gzip" + elif utils.is_lzma_alone_bytes(archive_bytes): + archive_type = "lzma" else: archive_type = "tar" return _extract_and_scan_archive( @@ -1024,7 +1030,9 @@ def security_scan( fail_fast=fail_fast, check_magic_bytes=check_magic_bytes, ) - elif pickle_bytes.startswith(b"\xfd7zXZ\x00"): + elif pickle_bytes.startswith(b"\xfd7zXZ\x00") or utils.is_lzma_alone_bytes( + pickle_bytes + ): return _extract_and_scan_archive( pickle_bytes, "lzma", @@ -1249,9 +1257,13 @@ def _security_scan_internal( stream.seek(0) 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") - ) or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar") + is_archive = ( + header_bytes.startswith( + (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") + ) + or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar") + or utils.is_lzma_alone_bytes(header_bytes) + ) if not is_archive: start_offset = utils.find_pickle_start_offset(stream) else: @@ -1413,9 +1425,13 @@ def _scan_and_load( current_pos = pickle_file.tell() 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") - ) or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar") + is_archive = ( + header_bytes.startswith( + (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") + ) + or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar") + or utils.is_lzma_alone_bytes(header_bytes) + ) if not is_archive: start_offset = utils.find_pickle_start_offset(header_bytes) else: @@ -1433,9 +1449,13 @@ def _scan_and_load( if not isinstance(pickle_file_or_bytes, bytes): 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") - ) or (len(data_bytes) >= 262 and data_bytes[257:262] == b"ustar") + is_archive = ( + data_bytes.startswith( + (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") + ) + or (len(data_bytes) >= 262 and data_bytes[257:262] == b"ustar") + or utils.is_lzma_alone_bytes(data_bytes) + ) if not is_archive: start_offset = utils.find_pickle_start_offset(data_bytes) else: