1#!/usr/bin/env python 2# -*- coding: utf-8 -*- 3 4# Copyright (c) 2021 Huawei Device Co., Ltd. 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16import bisect 17import copy 18import os 19import struct 20import tempfile 21from hashlib import sha256 22 23from log_exception import UPDATE_LOGGER 24from blocks_manager import BlocksManager 25from utils import SPARSE_IMAGE_MAGIC 26from utils import HEADER_INFO_FORMAT 27from utils import HEADER_INFO_LEN 28from utils import EXTEND_VALUE 29from utils import FILE_MAP_ZERO_KEY 30from utils import FILE_MAP_NONZERO_KEY 31from utils import FILE_MAP_COPY_KEY 32from utils import MAX_BLOCKS_PER_GROUP 33 34 35class FullUpdateImage: 36 """ 37 Full image processing class 38 """ 39 40 def __init__(self, target_package_images_dir, full_img_list, verse_script, 41 full_image_path_list, no_zip=False): 42 self.__target_package_images_dir = target_package_images_dir 43 self.__full_img_list = full_img_list 44 self.__verse_script = verse_script 45 self.__full_image_path_list = full_image_path_list 46 self.__no_zip = no_zip 47 48 def update_full_image(self): 49 """ 50 Processing of the full image 51 :return full_image_content_len_list: full image content length list 52 :return full_image_file_obj_list: full image temporary file list 53 """ 54 full_image_file_obj_list = [] 55 full_image_content_len_list = [] 56 for idx, each_name in enumerate(self.__full_img_list): 57 full_image_content = self.get_full_image_content( 58 self.__full_image_path_list[idx]) 59 if full_image_content is False: 60 UPDATE_LOGGER.print_log( 61 "Get full image content failed!", 62 log_type=UPDATE_LOGGER.ERROR_LOG) 63 return False, False 64 each_img = tempfile.NamedTemporaryFile( 65 prefix="full_image%s" % each_name, mode='wb') 66 each_img.write(full_image_content) 67 each_img.seek(0) 68 full_image_content_len_list.append(len(full_image_content)) 69 full_image_file_obj_list.append(each_img) 70 UPDATE_LOGGER.print_log( 71 "Image %s full processing completed" % each_name) 72 if not self.__no_zip: 73 # No zip mode (no script command) 74 if is_sparse_image(each_img.name): 75 sparse_image_write_cmd = \ 76 self.__verse_script.sparse_image_write(each_name) 77 cmd = '%s_WRITE_FLAG%s' % ( 78 each_name, sparse_image_write_cmd) 79 else: 80 raw_image_write_cmd = \ 81 self.__verse_script.raw_image_write( 82 each_name, each_name) 83 cmd = '%s_WRITE_FLAG%s' % ( 84 each_name, raw_image_write_cmd) 85 if each_name not in ("boot", "updater_boot", 86 "updater", "updater_b"): 87 self.__verse_script.add_command( 88 cmd=cmd) 89 90 UPDATE_LOGGER.print_log( 91 "All full image processing completed! image count: %d" % 92 len(self.__full_img_list)) 93 return full_image_content_len_list, full_image_file_obj_list 94 95 @staticmethod 96 def get_full_image_content(each_name): 97 """ 98 Obtain the full image content. 99 :param each_name: image name 100 :return content: full image content if available; false otherwise 101 """ 102 each_image_path = each_name 103 if not os.path.exists(each_image_path): 104 UPDATE_LOGGER.print_log( 105 "The file is missing " 106 "from the target package, " 107 "the component: %s cannot be full update processed. " % 108 each_image_path) 109 return False 110 with open(each_image_path, 'rb') as f_r: 111 content = f_r.read() 112 return content 113 114 115def is_sparse_image(img_path): 116 """ 117 Check whether the image is a sparse image. 118 :param img_path: image path 119 :return: 120 """ 121 with open(img_path, 'rb') as f_r: 122 image_content = f_r.read(HEADER_INFO_LEN) 123 try: 124 header_info = struct.unpack(HEADER_INFO_FORMAT, image_content) 125 except struct.error: 126 return False 127 is_sparse = IncUpdateImage.image_header_info_check(header_info)[-1] 128 if is_sparse: 129 UPDATE_LOGGER.print_log("Sparse image is not supported!") 130 raise RuntimeError 131 return is_sparse 132 133 134class IncUpdateImage: 135 """ 136 Increment update image class 137 """ 138 139 def __init__(self, image_path, map_path): 140 """ 141 Initialize the inc image. 142 :param image_path: img file path 143 :param map_path: map file path 144 """ 145 self.image_path = image_path 146 self.offset_value_list = [] 147 self.care_block_range = None 148 self.extended_range = None 149 self.reserved_blocks = BlocksManager("0") 150 self.file_map = [] 151 self.offset_index = [] 152 self.block_size = None 153 self.total_blocks = None 154 self.parse_sparse_image_file(image_path, map_path) 155 156 def parse_sparse_image_file(self, image_path, map_path): 157 """ 158 Parse the .img file. 159 :param image_path: img file path 160 :param map_path: map file path 161 """ 162 self.block_size = block_size = 4096 163 self.total_blocks = total_blocks = \ 164 os.path.getsize(self.image_path) // self.block_size 165 reference = b'\0' * self.block_size 166 with open(image_path, 'rb') as f_r: 167 care_value_list, offset_value_list = [], [] 168 nonzero_blocks = [] 169 for i in range(self.total_blocks): 170 blocks_data = f_r.read(self.block_size) 171 if blocks_data != reference: 172 nonzero_blocks.append(i) 173 nonzero_blocks.append(i + 1) 174 self.care_block_range = BlocksManager(nonzero_blocks) 175 care_value_list = list(self.care_block_range.range_data) 176 for idx, value in enumerate(care_value_list): 177 if idx != 0 and (idx + 1) % 2 == 0: 178 be_value = int(care_value_list[idx - 1]) 179 af_value = int(care_value_list[idx]) 180 file_tell = be_value * block_size 181 offset_value_list.append( 182 (be_value, af_value - be_value, 183 file_tell, None)) 184 185 self.offset_index = [i[0] for i in offset_value_list] 186 self.offset_value_list = offset_value_list 187 extended_range = \ 188 self.care_block_range.extend_value_to_blocks(EXTEND_VALUE) 189 all_blocks = BlocksManager(range_data=(0, total_blocks)) 190 self.extended_range = \ 191 extended_range.get_intersect_with_other(all_blocks). \ 192 get_subtract_with_other(self.care_block_range) 193 self.parse_block_map_file(map_path, f_r) 194 195 def parse_block_map_file(self, map_path, image_file_r): 196 """ 197 Parses the map file for blocks where files are contained in the image. 198 :param map_path: map file path 199 :param image_file_r: file reading object 200 :return: 201 """ 202 remain_range = self.care_block_range 203 temp_file_map = {} 204 205 with open(map_path, 'r') as f_r: 206 # Read the .map file and process each line. 207 for each_line in f_r.readlines(): 208 each_map_path, ranges_value = each_line.split(None, 1) 209 each_range = BlocksManager(ranges_value) 210 temp_file_map[each_map_path] = each_range 211 # each_range is contained in the remain range. 212 if each_range.size() != each_range. \ 213 get_intersect_with_other(remain_range).size(): 214 raise RuntimeError 215 # After the processing is complete, 216 # remove each_range from remain_range. 217 remain_range = remain_range.get_subtract_with_other(each_range) 218 reserved_blocks = self.reserved_blocks 219 # Remove reserved blocks from all blocks. 220 remain_range = remain_range.get_subtract_with_other(reserved_blocks) 221 222 # Divide all blocks into zero_blocks 223 # (if there are many) and nonzero_blocks. 224 zero_blocks_list = [] 225 nonzero_blocks_list = [] 226 nonzero_groups_list = [] 227 default_zero_block = ('\0' * self.block_size).encode() 228 229 nonzero_blocks_list, nonzero_groups_list, zero_blocks_list = \ 230 self.apply_remain_range( 231 default_zero_block, image_file_r, nonzero_blocks_list, 232 nonzero_groups_list, remain_range, zero_blocks_list) 233 234 temp_file_map = self.get_file_map( 235 nonzero_blocks_list, nonzero_groups_list, 236 reserved_blocks, temp_file_map, zero_blocks_list) 237 self.file_map = temp_file_map 238 239 def apply_remain_range(self, *args): 240 """ 241 Implement traversal processing of remain_range. 242 """ 243 default_zero_block, image_file_r, \ 244 nonzero_blocks_list, nonzero_groups_list, \ 245 remain_range, zero_blocks_list = args 246 for start_value, end_value in remain_range: 247 for each_value in range(start_value, end_value): 248 # bisect 二分查找，b在self.offset_index中的位置 249 idx = bisect.bisect_right(self.offset_index, each_value) - 1 250 chunk_start, _, file_pos, fill_data = \ 251 self.offset_value_list[idx] 252 data = self.get_file_data(self.block_size, chunk_start, 253 default_zero_block, each_value, 254 file_pos, fill_data, image_file_r) 255 256 zero_blocks_list, nonzero_blocks_list, nonzero_groups_list = \ 257 self.get_zero_nonzero_blocks_list( 258 data, default_zero_block, each_value, 259 nonzero_blocks_list, nonzero_groups_list, 260 zero_blocks_list) 261 return nonzero_blocks_list, nonzero_groups_list, zero_blocks_list 262 263 @staticmethod 264 def get_file_map(*args): 265 """ 266 Obtain the file map. 267 nonzero_blocks_list nonzero blocks list, 268 nonzero_groups_list nonzero groups list, 269 reserved_blocks reserved blocks , 270 temp_file_map temporary file map, 271 zero_blocks_list zero block list 272 :return temp_file_map file map 273 """ 274 nonzero_blocks_list, nonzero_groups_list, \ 275 reserved_blocks, temp_file_map, zero_blocks_list = args 276 if nonzero_blocks_list: 277 nonzero_groups_list.append(nonzero_blocks_list) 278 if zero_blocks_list: 279 temp_file_map[FILE_MAP_ZERO_KEY] = \ 280 BlocksManager(range_data=zero_blocks_list) 281 if nonzero_groups_list: 282 for i, blocks in enumerate(nonzero_groups_list): 283 temp_file_map["%s-%d" % (FILE_MAP_NONZERO_KEY, i)] = \ 284 BlocksManager(range_data=blocks) 285 if reserved_blocks: 286 temp_file_map[FILE_MAP_COPY_KEY] = reserved_blocks 287 return temp_file_map 288 289 @staticmethod 290 def get_zero_nonzero_blocks_list(*args): 291 """ 292 Get zero_blocks_list, nonzero_blocks_list, and nonzero_groups_list. 293 data: block data, 294 default_zero_block: default to zero block, 295 each_value: each value, 296 nonzero_blocks_list: nonzero_blocks_list, 297 nonzero_groups_list: nonzero_groups_list, 298 zero_blocks_list: zero_blocks_list, 299 :return new_zero_blocks_list: new zero blocks list, 300 :return new_nonzero_blocks_list: new nonzero blocks list, 301 :return new_nonzero_groups_list: new nonzero groups list. 302 """ 303 data, default_zero_block, each_value, \ 304 nonzero_blocks_list, nonzero_groups_list, \ 305 zero_blocks_list = args 306 # Check whether the data block is equal to the default zero_blocks. 307 if data == default_zero_block: 308 zero_blocks_list.append(each_value) 309 zero_blocks_list.append(each_value + 1) 310 else: 311 nonzero_blocks_list.append(each_value) 312 nonzero_blocks_list.append(each_value + 1) 313 # The number of nonzero_blocks is greater than 314 # or equal to the upper limit. 315 if len(nonzero_blocks_list) >= MAX_BLOCKS_PER_GROUP: 316 nonzero_groups_list.append(nonzero_blocks_list) 317 nonzero_blocks_list = [] 318 new_zero_blocks_list, new_nonzero_blocks_list, \ 319 new_nonzero_groups_list = \ 320 copy.copy(zero_blocks_list), \ 321 copy.copy(nonzero_blocks_list),\ 322 copy.copy(nonzero_groups_list) 323 return new_zero_blocks_list, new_nonzero_blocks_list, \ 324 new_nonzero_groups_list 325 326 @staticmethod 327 def get_file_data(*args): 328 """ 329 Get the file data. 330 block_size: blocksize, 331 chunk_start: the start position of chunk, 332 default_zero_block: default to zero blocks, 333 each_value: each_value, 334 file_pos: file position, 335 fill_data: data, 336 image_file_r: read file object, 337 :return data: Get the file data. 338 """ 339 block_size, chunk_start, default_zero_block, each_value, \ 340 file_pos, fill_data, image_file_r = args 341 if file_pos is not None: 342 file_pos += (each_value - chunk_start) * block_size 343 image_file_r.seek(file_pos, os.SEEK_SET) 344 data = image_file_r.read(block_size) 345 else: 346 if fill_data == default_zero_block[:4]: 347 data = default_zero_block 348 else: 349 data = None 350 return data 351 352 def range_sha256(self, ranges): 353 """ 354 range sha256 hash content 355 :param ranges: ranges value 356 :return: 357 """ 358 hash_obj = sha256() 359 for data in self.__get_blocks_set_data(ranges): 360 hash_obj.update(data) 361 return hash_obj.hexdigest() 362 363 def write_range_data_2_fd(self, ranges, file_obj): 364 """ 365 write range data to fd 366 :param ranges: ranges obj 367 :param file_obj: file obj 368 :return: 369 """ 370 for data in self.__get_blocks_set_data(ranges): 371 file_obj.write(data) 372 373 def get_ranges(self, ranges): 374 """ 375 get ranges value 376 :param ranges: ranges 377 :return: ranges value 378 """ 379 return [each_data for each_data in self.__get_blocks_set_data(ranges)] 380 381 def __get_blocks_set_data(self, blocks_set_data): 382 """ 383 Get the range data. 384 """ 385 with open(self.image_path, 'rb') as f_r: 386 for start, end in blocks_set_data: 387 diff_value = end - start 388 idx = bisect.bisect_right(self.offset_index, start) - 1 389 chunk_start, chunk_len, file_pos, fill_data = \ 390 self.offset_value_list[idx] 391 392 remain = chunk_len - (start - chunk_start) 393 this_read = min(remain, diff_value) 394 if file_pos is not None: 395 pos = file_pos + ((start - chunk_start) * self.block_size) 396 f_r.seek(pos, os.SEEK_SET) 397 yield f_r.read(this_read * self.block_size) 398 else: 399 yield fill_data * (this_read * (self.block_size >> 2)) 400 diff_value -= this_read 401 402 while diff_value > 0: 403 idx += 1 404 chunk_start, chunk_len, file_pos, fill_data = \ 405 self.offset_value_list[idx] 406 this_read = min(chunk_len, diff_value) 407 if file_pos is not None: 408 f_r.seek(file_pos, os.SEEK_SET) 409 yield f_r.read(this_read * self.block_size) 410 else: 411 yield fill_data * (this_read * (self.block_size >> 2)) 412 diff_value -= this_read 413 414 @staticmethod 415 def image_header_info_check(header_info): 416 """ 417 Check for new messages of the header_info image. 418 :param header_info: header_info 419 :return: 420 """ 421 image_flag = True 422 # Sparse mirroring header ID. The magic value is fixed to 0xED26FF3A. 423 magic_info = header_info[0] 424 # major version number 425 major_version = header_info[1] 426 # minor version number 427 minor_version = header_info[2] 428 # Length of the header information. 429 # The value is fixed to 28 characters. 430 header_info_size = header_info[3] 431 # Header information size of the chunk. 432 # The length is fixed to 12 characters. 433 chunk_header_info_size = header_info[4] 434 # Number of bytes of a block. The default size is 4096. 435 block_size = header_info[5] 436 # Total number of blocks contained in the current image 437 # (number of blocks in a non-sparse image) 438 total_blocks = header_info[6] 439 # Total number of chunks contained in the current image 440 total_chunks = header_info[7] 441 if magic_info != SPARSE_IMAGE_MAGIC: 442 UPDATE_LOGGER.print_log( 443 "SparseImage head Magic should be 0xED26FF3A!") 444 image_flag = False 445 if major_version != 1 or minor_version != 0: 446 UPDATE_LOGGER.print_log( 447 "SparseImage Only supported major version with " 448 "minor version 1.0!") 449 image_flag = False 450 if header_info_size != 28: 451 UPDATE_LOGGER.print_log( 452 "SparseImage header info size must be 28! size: %u." % 453 header_info_size) 454 image_flag = False 455 if chunk_header_info_size != 12: 456 UPDATE_LOGGER.print_log( 457 "SparseImage Chunk header size mast to be 12! size: %u." % 458 chunk_header_info_size) 459 image_flag = False 460 ret_args = [block_size, chunk_header_info_size, header_info_size, 461 magic_info, total_blocks, total_chunks, image_flag] 462 return ret_args 463