sparse_img.py revision 5fcaaeffc3d132f1223adf7ee8230b6815477af5
1# Copyright (C) 2014 The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15import bisect 16import os 17import struct 18from hashlib import sha1 19 20import rangelib 21 22 23class SparseImage(object): 24 """Wraps a sparse image file into an image object. 25 26 Wraps a sparse image file (and optional file map and clobbered_blocks) into 27 an image object suitable for passing to BlockImageDiff. file_map contains 28 the mapping between files and their blocks. clobbered_blocks contains the set 29 of blocks that should be always written to the target regardless of the old 30 contents (i.e. copying instead of patching). clobbered_blocks should be in 31 the form of a string like "0" or "0 1-5 8". 32 """ 33 34 def __init__(self, simg_fn, file_map_fn=None, clobbered_blocks=None): 35 self.simg_f = f = open(simg_fn, "rb") 36 37 header_bin = f.read(28) 38 header = struct.unpack("<I4H4I", header_bin) 39 40 magic = header[0] 41 major_version = header[1] 42 minor_version = header[2] 43 file_hdr_sz = header[3] 44 chunk_hdr_sz = header[4] 45 self.blocksize = blk_sz = header[5] 46 self.total_blocks = total_blks = header[6] 47 total_chunks = header[7] 48 49 if magic != 0xED26FF3A: 50 raise ValueError("Magic should be 0xED26FF3A but is 0x%08X" % (magic,)) 51 if major_version != 1 or minor_version != 0: 52 raise ValueError("I know about version 1.0, but this is version %u.%u" % 53 (major_version, minor_version)) 54 if file_hdr_sz != 28: 55 raise ValueError("File header size was expected to be 28, but is %u." % 56 (file_hdr_sz,)) 57 if chunk_hdr_sz != 12: 58 raise ValueError("Chunk header size was expected to be 12, but is %u." % 59 (chunk_hdr_sz,)) 60 61 print("Total of %u %u-byte output blocks in %u input chunks." 62 % (total_blks, blk_sz, total_chunks)) 63 64 pos = 0 # in blocks 65 care_data = [] 66 self.offset_map = offset_map = [] 67 self.clobbered_blocks = rangelib.RangeSet(data=clobbered_blocks) 68 69 for i in range(total_chunks): 70 header_bin = f.read(12) 71 header = struct.unpack("<2H2I", header_bin) 72 chunk_type = header[0] 73 chunk_sz = header[2] 74 total_sz = header[3] 75 data_sz = total_sz - 12 76 77 if chunk_type == 0xCAC1: 78 if data_sz != (chunk_sz * blk_sz): 79 raise ValueError( 80 "Raw chunk input size (%u) does not match output size (%u)" % 81 (data_sz, chunk_sz * blk_sz)) 82 else: 83 care_data.append(pos) 84 care_data.append(pos + chunk_sz) 85 offset_map.append((pos, chunk_sz, f.tell(), None)) 86 pos += chunk_sz 87 f.seek(data_sz, os.SEEK_CUR) 88 89 elif chunk_type == 0xCAC2: 90 fill_data = f.read(4) 91 care_data.append(pos) 92 care_data.append(pos + chunk_sz) 93 offset_map.append((pos, chunk_sz, None, fill_data)) 94 pos += chunk_sz 95 96 elif chunk_type == 0xCAC3: 97 if data_sz != 0: 98 raise ValueError("Don't care chunk input size is non-zero (%u)" % 99 (data_sz)) 100 else: 101 pos += chunk_sz 102 103 elif chunk_type == 0xCAC4: 104 raise ValueError("CRC32 chunks are not supported") 105 106 else: 107 raise ValueError("Unknown chunk type 0x%04X not supported" % 108 (chunk_type,)) 109 110 self.care_map = rangelib.RangeSet(care_data) 111 self.offset_index = [i[0] for i in offset_map] 112 113 if file_map_fn: 114 self.LoadFileBlockMap(file_map_fn, self.clobbered_blocks) 115 else: 116 self.file_map = {"__DATA": self.care_map} 117 118 def ReadRangeSet(self, ranges): 119 return [d for d in self._GetRangeData(ranges)] 120 121 def TotalSha1(self, include_clobbered_blocks=False): 122 """Return the SHA-1 hash of all data in the 'care' regions. 123 124 If include_clobbered_blocks is True, it returns the hash including the 125 clobbered_blocks.""" 126 ranges = self.care_map 127 if not include_clobbered_blocks: 128 ranges.subtract(self.clobbered_blocks) 129 h = sha1() 130 for d in self._GetRangeData(ranges): 131 h.update(d) 132 return h.hexdigest() 133 134 def _GetRangeData(self, ranges): 135 """Generator that produces all the image data in 'ranges'. The 136 number of individual pieces returned is arbitrary (and in 137 particular is not necessarily equal to the number of ranges in 138 'ranges'. 139 140 This generator is stateful -- it depends on the open file object 141 contained in this SparseImage, so you should not try to run two 142 instances of this generator on the same object simultaneously.""" 143 144 f = self.simg_f 145 for s, e in ranges: 146 to_read = e-s 147 idx = bisect.bisect_right(self.offset_index, s) - 1 148 chunk_start, chunk_len, filepos, fill_data = self.offset_map[idx] 149 150 # for the first chunk we may be starting partway through it. 151 remain = chunk_len - (s - chunk_start) 152 this_read = min(remain, to_read) 153 if filepos is not None: 154 p = filepos + ((s - chunk_start) * self.blocksize) 155 f.seek(p, os.SEEK_SET) 156 yield f.read(this_read * self.blocksize) 157 else: 158 yield fill_data * (this_read * (self.blocksize >> 2)) 159 to_read -= this_read 160 161 while to_read > 0: 162 # continue with following chunks if this range spans multiple chunks. 163 idx += 1 164 chunk_start, chunk_len, filepos, fill_data = self.offset_map[idx] 165 this_read = min(chunk_len, to_read) 166 if filepos is not None: 167 f.seek(filepos, os.SEEK_SET) 168 yield f.read(this_read * self.blocksize) 169 else: 170 yield fill_data * (this_read * (self.blocksize >> 2)) 171 to_read -= this_read 172 173 def LoadFileBlockMap(self, fn, clobbered_blocks): 174 remaining = self.care_map 175 self.file_map = out = {} 176 177 with open(fn) as f: 178 for line in f: 179 fn, ranges = line.split(None, 1) 180 ranges = rangelib.RangeSet.parse(ranges) 181 out[fn] = ranges 182 assert ranges.size() == ranges.intersect(remaining).size() 183 184 # Currently we assume that blocks in clobbered_blocks are not part of 185 # any file. 186 assert not clobbered_blocks.overlaps(ranges) 187 remaining = remaining.subtract(ranges) 188 189 remaining = remaining.subtract(clobbered_blocks) 190 191 # For all the remaining blocks in the care_map (ie, those that 192 # aren't part of the data for any file nor part of the clobbered_blocks), 193 # divide them into blocks that are all zero and blocks that aren't. 194 # (Zero blocks are handled specially because (1) there are usually 195 # a lot of them and (2) bsdiff handles files with long sequences of 196 # repeated bytes especially poorly.) 197 198 zero_blocks = [] 199 nonzero_blocks = [] 200 reference = '\0' * self.blocksize 201 202 f = self.simg_f 203 for s, e in remaining: 204 for b in range(s, e): 205 idx = bisect.bisect_right(self.offset_index, b) - 1 206 chunk_start, _, filepos, fill_data = self.offset_map[idx] 207 if filepos is not None: 208 filepos += (b-chunk_start) * self.blocksize 209 f.seek(filepos, os.SEEK_SET) 210 data = f.read(self.blocksize) 211 else: 212 if fill_data == reference[:4]: # fill with all zeros 213 data = reference 214 else: 215 data = None 216 217 if data == reference: 218 zero_blocks.append(b) 219 zero_blocks.append(b+1) 220 else: 221 nonzero_blocks.append(b) 222 nonzero_blocks.append(b+1) 223 224 out["__ZERO"] = rangelib.RangeSet(data=zero_blocks) 225 out["__NONZERO"] = rangelib.RangeSet(data=nonzero_blocks) 226 out["__COPY"] = clobbered_blocks 227 228 def ResetFileMap(self): 229 """Throw away the file map and treat the entire image as 230 undifferentiated data.""" 231 self.file_map = {"__DATA": self.care_map} 232