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

1import itertools 

2import lzma 

3import re 

4import zlib 

5from collections.abc import Callable 

6from pathlib import Path 

7 

8import zstandard 

9 

10from unblob.file_utils import ( 

11 Endian, 

12 FileSystem, 

13 InvalidInputFormat, 

14 StructParser, 

15 iterate_file, 

16) 

17from unblob.models import ( 

18 Extractor, 

19 ExtractResult, 

20 File, 

21 HandlerDoc, 

22 HandlerType, 

23 Reference, 

24 Regex, 

25 StructHandler, 

26 ValidChunk, 

27) 

28 

29# [Ref] https://github.com/freebsd/freebsd-src/tree/master/sys/geom/uzip 

30C_DEFINITIONS = r""" 

31 typedef struct uzip_header{ 

32 char magic[16]; 

33 char format[112]; 

34 uint32_t block_size; 

35 uint32_t block_count; 

36 uint64_t toc[block_count]; 

37 } uzip_header_t; 

38""" 

39 

40HEADER_STRUCT = "uzip_header_t" 

41 

42ZLIB_COMPRESSION = "#!/bin/sh\x0a#V2.0\x20" 

43LZMA_COMPRESSION = "#!/bin/sh\x0a#L3.0\x0a" 

44ZSTD_COMPRESSION = "#!/bin/sh\x0a#Z4.0\x20" 

45 

46 

47class Decompressor: 

48 DECOMPRESSOR: Callable 

49 

50 def __init__(self): 

51 self._decompressor = self.DECOMPRESSOR() 

52 

53 def decompress(self, data: bytes) -> bytes: 

54 return self._decompressor.decompress(data) 

55 

56 def flush(self) -> bytes: 

57 return b"" 

58 

59 

60class LZMADecompressor(Decompressor): 

61 DECOMPRESSOR = lzma.LZMADecompressor 

62 

63 

64class ZLIBDecompressor(Decompressor): 

65 DECOMPRESSOR = zlib.decompressobj 

66 

67 def flush(self) -> bytes: 

68 return self._decompressor.flush() 

69 

70 

71class ZSTDDecompressor(Decompressor): 

72 DECOMPRESSOR = zstandard.ZstdDecompressor 

73 

74 

75DECOMPRESS_METHOD: dict[bytes, type[Decompressor]] = { 

76 ZLIB_COMPRESSION.encode(): ZLIBDecompressor, 

77 LZMA_COMPRESSION.encode(): LZMADecompressor, 

78 ZSTD_COMPRESSION.encode(): ZSTDDecompressor, 

79} 

80 

81 

82def is_toc_sorted(toc) -> bool: 

83 return all(current <= nxt for current, nxt in itertools.pairwise(toc)) 

84 

85 

86class UZIPExtractor(Extractor): 

87 def extract(self, inpath: Path, outdir: Path): 

88 with File.from_path(inpath) as infile: 

89 parser = StructParser(C_DEFINITIONS) 

90 header = parser.parse(HEADER_STRUCT, infile, Endian.BIG) 

91 fs = FileSystem(outdir) 

92 outpath = Path(inpath.stem) 

93 

94 try: 

95 decompressor_cls = DECOMPRESS_METHOD[header.magic] 

96 except LookupError: 

97 raise InvalidInputFormat("unsupported compression format") from None 

98 

99 with fs.open(outpath, "wb+") as outfile: 

100 for current_offset, next_offset in zip( 

101 header.toc[:-1], header.toc[1:], strict=False 

102 ): 

103 compressed_len = next_offset - current_offset 

104 if compressed_len == 0: 

105 continue 

106 decompressor = decompressor_cls() 

107 for chunk in iterate_file(infile, current_offset, compressed_len): 

108 outfile.write(decompressor.decompress(chunk)) 

109 outfile.write(decompressor.flush()) 

110 return ExtractResult(reports=fs.problems) 

111 

112 

113class UZIPHandler(StructHandler): 

114 NAME = "uzip" 

115 PATTERNS = [ 

116 Regex(re.escape(ZLIB_COMPRESSION)), 

117 Regex(re.escape(LZMA_COMPRESSION)), 

118 Regex(re.escape(ZSTD_COMPRESSION)), 

119 ] 

120 HEADER_STRUCT = HEADER_STRUCT 

121 C_DEFINITIONS = C_DEFINITIONS 

122 EXTRACTOR = UZIPExtractor() 

123 

124 DOC = HandlerDoc( 

125 name="UZIP", 

126 description="FreeBSD UZIP is a block-based compressed disk image format. It uses a table of contents to index compressed blocks, supporting ZLIB, LZMA, and ZSTD compression algorithms.", 

127 handler_type=HandlerType.COMPRESSION, 

128 vendor="FreeBSD", 

129 references=[ 

130 Reference( 

131 title="FreeBSD UZIP Documentation", 

132 url="https://github.com/freebsd/freebsd-src/tree/master/sys/geom/uzip", 

133 ), 

134 ], 

135 limitations=[], 

136 ) 

137 

138 def is_valid_header(self, header) -> bool: 

139 return ( 

140 header.block_count > 0 

141 and header.block_size > 0 

142 and header.block_size % 512 == 0 

143 and is_toc_sorted(header.toc) 

144 ) 

145 

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

147 header = self.parse_header(file, Endian.BIG) 

148 

149 if not self.is_valid_header(header): 

150 raise InvalidInputFormat("Invalid uzip header.") 

151 

152 # take the last TOC block offset, end of file is that block offset, 

153 # starting from the start offset 

154 end_offset = start_offset + header.toc[-1] 

155 return ValidChunk( 

156 start_offset=start_offset, 

157 end_offset=end_offset, 

158 )