Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/unblob/handlers/compression/lzfse.py: 42%

1import io 

2from enum import Enum 

3from pathlib import Path 

4 

5import lzfse 

6 

7from unblob.file_utils import ( 

8 Endian, 

9 File, 

10 FileSystem, 

11 InvalidInputFormat, 

12 StructParser, 

13) 

14from unblob.models import ( 

15 Extractor, 

16 ExtractResult, 

17 Handler, 

18 HandlerDoc, 

19 HandlerType, 

20 HexString, 

21 Reference, 

22 ValidChunk, 

23) 

24 

25 

26class LZFSEMagic(bytes, Enum): 

27 END = b"bvx$" # end-of-stream block 

28 UNCOMPRESSED = b"bvx-" # raw block 

29 LZVN = b"bvxn" # LZVN-compressed block 

30 LZFSE_V1 = b"bvx1" # LZFSE v1 block (legacy) 

31 LZFSE_V2 = b"bvx2" # LZFSE v2 block 

32 

33 

34# sizeof(lzfse_compressed_block_header_v1), including struct alignment padding 

35_V1_HEADER_SIZE = 772 

36 

37# length of every LZFSE block magic 

38MAGIC_LEN = 4 

39 

40# a stream decoding to zero bytes: an empty raw block followed by the end block, 

41# as emitted by lzfse.compress(b"") 

42_EMPTY_STREAM = LZFSEMagic.UNCOMPRESSED + b"\x00\x00\x00\x00" + LZFSEMagic.END 

43 

44C_DEFINITIONS = r""" 

45 typedef struct lzfse_uncompressed { 

46 char magic[4]; 

47 uint32 n_raw_bytes; 

48 } lzfse_uncompressed_t; 

49 

50 typedef struct lzvn_compressed { 

51 char magic[4]; 

52 uint32 n_raw_bytes; 

53 uint32 n_payload_bytes; 

54 } lzvn_compressed_t; 

55 

56 typedef struct lzfse_v1 { 

57 char magic[4]; 

58 uint32 n_raw_bytes; 

59 uint32 n_payload_bytes; 

60 uint32 n_literals; 

61 uint32 n_matches; 

62 uint32 n_literal_payload_bytes; 

63 uint32 n_lmd_payload_bytes; 

64 } lzfse_v1_t; 

65 

66 typedef struct lzfse_v2 { 

67 char magic[4]; 

68 uint32 n_raw_bytes; 

69 uint64 packed_fields_0; 

70 uint64 packed_fields_1; 

71 uint64 packed_fields_2; 

72 } lzfse_v2_t; 

73""" 

74 

75_parser = StructParser(C_DEFINITIONS) 

76 

77 

78class LZFSEExtractor(Extractor): 

79 def extract(self, inpath: Path, outdir: Path) -> ExtractResult | None: 

80 fs = FileSystem(outdir) 

81 with File.from_path(inpath) as file: 

82 content = file.read() 

83 # lzfse.decompress raises on a stream that decodes to zero bytes, 

84 # so write the empty output ourselves. 

85 decompressed = ( 

86 b"" if content == _EMPTY_STREAM else lzfse.decompress(content) 

87 ) 

88 fs.write_bytes(Path(f"{inpath.stem}.bin"), decompressed) 

89 return ExtractResult(reports=fs.problems) 

90 

91 

92class LZFSEHandler(Handler): 

93 NAME = "lzfse" 

94 

95 PATTERNS = [ 

96 HexString("62 76 78 2D"), # "bvx-" uncompressed block 

97 HexString("62 76 78 31"), # "bvx1" LZFSE v1 compressed block (legacy) 

98 HexString("62 76 78 6E"), # "bvxn" LZVN compressed block 

99 HexString("62 76 78 32"), # "bvx2" LZFSE v2 compressed block 

100 ] 

101 

102 EXTRACTOR = LZFSEExtractor() 

103 

104 DOC = HandlerDoc( 

105 name="LZFSE", 

106 description="LZFSE is a lossless compression algorithm developed by Apple and open-sourced in 2016. It combines Lempel-Ziv back-references with Finite State Entropy coding and is the default compression format used in iOS and macOS firmware images.", 

107 handler_type=HandlerType.COMPRESSION, 

108 vendor="Apple", 

109 references=[ 

110 Reference( 

111 title="lzfse - Apple open-source LZFSE library", 

112 url="https://github.com/lzfse/lzfse", 

113 ), 

114 ], 

115 limitations=[], 

116 ) 

117 

118 def calculate_chunk(self, file: File, start_offset: int) -> ValidChunk | None: 

119 # An LZFSE stream is a sequence of blocks terminated by an end-of-stream 

120 # block. Walk the blocks using each header's declared size instead of 

121 # scanning for "bvx$", which could otherwise be matched inside payload data. 

122 offset = start_offset 

123 magic = file[offset : offset + MAGIC_LEN] 

124 while magic != LZFSEMagic.END: 

125 if len(magic) < MAGIC_LEN: 

126 raise InvalidInputFormat("Truncated LZFSE stream: no end block") 

127 block_size = self._block_size(file, offset, magic) 

128 if block_size <= 0: 

129 raise InvalidInputFormat("Invalid LZFSE block size") 

130 offset += block_size 

131 magic = file[offset : offset + MAGIC_LEN] 

132 

133 return ValidChunk(start_offset=start_offset, end_offset=offset + MAGIC_LEN) 

134 

135 @staticmethod 

136 def _block_size(file: File, offset: int, magic: bytes) -> int: 

137 """Size in bytes of the LZFSE block at offset, including its header.""" 

138 file.seek(offset, io.SEEK_SET) 

139 match magic: 

140 case LZFSEMagic.UNCOMPRESSED: 

141 header = _parser.parse("lzfse_uncompressed_t", file, Endian.LITTLE) 

142 size = 8 + header.n_raw_bytes 

143 case LZFSEMagic.LZVN: 

144 header = _parser.parse("lzvn_compressed_t", file, Endian.LITTLE) 

145 size = 12 + header.n_payload_bytes 

146 case LZFSEMagic.LZFSE_V1: 

147 header = _parser.parse("lzfse_v1_t", file, Endian.LITTLE) 

148 size = ( 

149 _V1_HEADER_SIZE 

150 + header.n_literal_payload_bytes 

151 + header.n_lmd_payload_bytes 

152 ) 

153 case LZFSEMagic.LZFSE_V2: 

154 header = _parser.parse("lzfse_v2_t", file, Endian.LITTLE) 

155 # v2 packs the sizes into bit-fields across three uint64s. 

156 n_literal = (header.packed_fields_0 >> 20) & 0xFFFFF 

157 n_lmd = (header.packed_fields_1 >> 40) & 0xFFFFF 

158 header_size = header.packed_fields_2 & 0xFFFFFFFF 

159 size = header_size + n_literal + n_lmd 

160 case _: 

161 raise InvalidInputFormat(f"Unknown LZFSE block magic: {magic!r}") 

162 return size