1#!/usr/bin/env python3 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 OPTIONS_MANAGER 26from utils import EXTEND_VALUE 27from utils import FILE_MAP_ZERO_KEY 28from utils import FILE_MAP_NONZERO_KEY 29from utils import FILE_MAP_COPY_KEY 30from utils import MAX_BLOCKS_PER_GROUP 31from utils import FORBIDEN_UPDATE_IMAGE_LIST 32 33 34class FullUpdateImage: 35 """ 36 Full image processing class 37 """ 38 39 def __init__(self, target_package_images_dir, 40 full_img_list, full_img_name_list, 41 verse_script, full_image_path_list, 42 no_zip=False): 43 self.target_package_images_dir = target_package_images_dir 44 self.full_img_list = full_img_list 45 self.full_img_name_list = full_img_name_list 46 self.verse_script = verse_script 47 self.full_image_path_list = full_image_path_list 48 self.no_zip = no_zip 49 50 def update_full_image(self): 51 """ 52 Processing of the full image 53 :return full_image_content_len_list: full image content length list 54 :return full_image_file_obj_list: full image temporary file list 55 """ 56 full_image_file_obj_list = [] 57 full_image_content_len_list = [] 58 for idx, each_name in enumerate(self.full_img_list): 59 full_image_content = self.get_full_image_content( 60 self.full_image_path_list[idx]) 61 img_name = self.full_img_name_list[idx][:-4] 62 if full_image_content is False: 63 UPDATE_LOGGER.print_log( 64 "Get full image content failed!", 65 log_type=UPDATE_LOGGER.ERROR_LOG) 66 return False, False 67 each_img = tempfile.NamedTemporaryFile( 68 dir=self.target_package_images_dir, 69 prefix="full_image%s" % img_name, mode='wb') 70 each_img.write(full_image_content) 71 each_img.seek(0) 72 full_image_content_len_list.append(len(full_image_content)) 73 full_image_file_obj_list.append(each_img) 74 UPDATE_LOGGER.print_log( 75 "Image %s full processing completed" % img_name) 76 if not self.no_zip: 77 # No zip mode (no script command) 78 image_write_cmd = \ 79 self.verse_script.image_write(each_name, img_name, each_img.name) 80 cmd = '%s_WRITE_FLAG%s' % (each_name, image_write_cmd) 81 if each_name not in FORBIDEN_UPDATE_IMAGE_LIST: 82 self.verse_script.add_command(cmd=cmd) 83 84 UPDATE_LOGGER.print_log( 85 "All full image processing completed! image count: %d" % 86 len(self.full_img_list)) 87 return full_image_content_len_list, full_image_file_obj_list 88 89 @staticmethod 90 def get_full_image_content(each_name): 91 """ 92 Obtain the full image content. 93 :param each_name: image name 94 :return content: full image content if available; false otherwise 95 """ 96 each_image_path = each_name 97 if not os.path.exists(each_image_path): 98 UPDATE_LOGGER.print_log( 99 "The file is missing " 100 "from the target package, " 101 "the component: %s cannot be full update processed. " % 102 each_image_path) 103 return False 104 with open(each_image_path, 'rb') as f_r: 105 content = f_r.read() 106 return content 107 108 109class IncUpdateImage: 110 """ 111 Increment update image class 112 """ 113 114 def __init__(self, image_path, map_path): 115 """ 116 Initialize the inc image. 117 :param image_path: img file path 118 :param map_path: map file path 119 """ 120 self.image_path = image_path 121 self.map_path = map_path 122 self.offset_value_list = [] 123 self.care_block_range = None 124 self.extended_range = None 125 self.reserved_blocks = BlocksManager("0") 126 self.file_map = [] 127 self.offset_index = [] 128 self.block_size = None 129 self.total_blocks = None 130 self.parse_raw_image_file(image_path, map_path) 131 132 def parse_raw_image_file(self, image_path, map_path): 133 """ 134 Parse the .img file. 135 :param image_path: img file path 136 :param map_path: map file path 137 """ 138 self.block_size = block_size = 4096 139 self.total_blocks = total_blocks = \ 140 os.path.getsize(self.image_path) // self.block_size 141 reference = b'\0' * self.block_size 142 with open(image_path, 'rb') as f_r: 143 care_value_list, offset_value_list = [], [] 144 nonzero_blocks = [] 145 for i in range(self.total_blocks): 146 blocks_data = f_r.read(self.block_size) 147 if blocks_data != reference: 148 nonzero_blocks.append(i) 149 nonzero_blocks.append(i + 1) 150 self.care_block_range = BlocksManager(nonzero_blocks) 151 care_value_list = list(self.care_block_range.range_data) 152 for idx, value in enumerate(care_value_list): 153 if idx != 0 and (idx + 1) % 2 == 0: 154 be_value = int(care_value_list[idx - 1]) 155 af_value = int(care_value_list[idx]) 156 file_tell = be_value * block_size 157 offset_value_list.append( 158 (be_value, af_value - be_value, 159 file_tell, None)) 160 161 self.offset_index = [i[0] for i in offset_value_list] 162 self.offset_value_list = offset_value_list 163 extended_range = \ 164 self.care_block_range.extend_value_to_blocks(EXTEND_VALUE) 165 all_blocks = BlocksManager(range_data=(0, total_blocks)) 166 self.extended_range = \ 167 extended_range.get_intersect_with_other(all_blocks). \ 168 get_subtract_with_other(self.care_block_range) 169 self.parse_block_map_file(map_path, f_r) 170 171 def parse_block_map_file(self, map_path, image_file_r): 172 """ 173 Parses the map file for blocks where files are contained in the image. 174 :param map_path: map file path 175 :param image_file_r: file reading object 176 :return: 177 """ 178 remain_range = self.care_block_range 179 temp_file_map = {} 180 181 with open(map_path, 'r') as f_r: 182 # Read the .map file and process each line. 183 for each_line in f_r.readlines(): 184 each_map_path, ranges_value = each_line.split(None, 1) 185 each_range = BlocksManager(ranges_value) 186 temp_file_map[each_map_path] = each_range 187 # each_range is contained in the remain range. 188 if each_range.size() != each_range. \ 189 get_intersect_with_other(remain_range).size(): 190 raise RuntimeError 191 # After the processing is complete, 192 # remove each_range from remain_range. 193 remain_range = remain_range.get_subtract_with_other(each_range) 194 reserved_blocks = self.reserved_blocks 195 # Remove reserved blocks from all blocks. 196 remain_range = remain_range.get_subtract_with_other(reserved_blocks) 197 198 # Divide all blocks into zero_blocks 199 # (if there are many) and nonzero_blocks. 200 zero_blocks_list = [] 201 nonzero_blocks_list = [] 202 nonzero_groups_list = [] 203 default_zero_block = ('\0' * self.block_size).encode() 204 205 nonzero_blocks_list, nonzero_groups_list, zero_blocks_list = \ 206 self.apply_remain_range( 207 default_zero_block, image_file_r, nonzero_blocks_list, 208 nonzero_groups_list, remain_range, zero_blocks_list) 209 210 temp_file_map = self.get_file_map( 211 nonzero_blocks_list, nonzero_groups_list, 212 reserved_blocks, temp_file_map, zero_blocks_list) 213 self.file_map = temp_file_map 214 215 def apply_remain_range(self, *args): 216 """ 217 Implement traversal processing of remain_range. 218 """ 219 default_zero_block, image_file_r, \ 220 nonzero_blocks_list, nonzero_groups_list, \ 221 remain_range, zero_blocks_list = args 222 for start_value, end_value in remain_range: 223 for each_value in range(start_value, end_value): 224 # bisect 二分查找，b在self.offset_index中的位置 225 idx = bisect.bisect_right(self.offset_index, each_value) - 1 226 chunk_start, _, file_pos, fill_data = \ 227 self.offset_value_list[idx] 228 data = self.get_file_data(self.block_size, chunk_start, 229 default_zero_block, each_value, 230 file_pos, fill_data, image_file_r) 231 232 zero_blocks_list, nonzero_blocks_list, nonzero_groups_list = \ 233 self.get_zero_nonzero_blocks_list( 234 data, default_zero_block, each_value, 235 nonzero_blocks_list, nonzero_groups_list, 236 zero_blocks_list) 237 return nonzero_blocks_list, nonzero_groups_list, zero_blocks_list 238 239 @staticmethod 240 def get_file_map(*args): 241 """ 242 Obtain the file map. 243 nonzero_blocks_list nonzero blocks list, 244 nonzero_groups_list nonzero groups list, 245 reserved_blocks reserved blocks , 246 temp_file_map temporary file map, 247 zero_blocks_list zero block list 248 :return temp_file_map file map 249 """ 250 nonzero_blocks_list, nonzero_groups_list, \ 251 reserved_blocks, temp_file_map, zero_blocks_list = args 252 if nonzero_blocks_list: 253 nonzero_groups_list.append(nonzero_blocks_list) 254 if zero_blocks_list: 255 temp_file_map[FILE_MAP_ZERO_KEY] = \ 256 BlocksManager(range_data=zero_blocks_list) 257 if nonzero_groups_list: 258 for i, blocks in enumerate(nonzero_groups_list): 259 temp_file_map["%s-%d" % (FILE_MAP_NONZERO_KEY, i)] = \ 260 BlocksManager(range_data=blocks) 261 if reserved_blocks: 262 temp_file_map[FILE_MAP_COPY_KEY] = reserved_blocks 263 return temp_file_map 264 265 @staticmethod 266 def get_zero_nonzero_blocks_list(*args): 267 """ 268 Get zero_blocks_list, nonzero_blocks_list, and nonzero_groups_list. 269 data: block data, 270 default_zero_block: default to zero block, 271 each_value: each value, 272 nonzero_blocks_list: nonzero_blocks_list, 273 nonzero_groups_list: nonzero_groups_list, 274 zero_blocks_list: zero_blocks_list, 275 :return new_zero_blocks_list: new zero blocks list, 276 :return new_nonzero_blocks_list: new nonzero blocks list, 277 :return new_nonzero_groups_list: new nonzero groups list. 278 """ 279 data, default_zero_block, each_value, \ 280 nonzero_blocks_list, nonzero_groups_list, \ 281 zero_blocks_list = args 282 # Check whether the data block is equal to the default zero_blocks. 283 if data == default_zero_block: 284 zero_blocks_list.append(each_value) 285 zero_blocks_list.append(each_value + 1) 286 else: 287 nonzero_blocks_list.append(each_value) 288 nonzero_blocks_list.append(each_value + 1) 289 # The number of nonzero_blocks is greater than 290 # or equal to the upper limit. 291 if len(nonzero_blocks_list) >= MAX_BLOCKS_PER_GROUP: 292 nonzero_groups_list.append(nonzero_blocks_list) 293 nonzero_blocks_list = [] 294 new_zero_blocks_list, new_nonzero_blocks_list, \ 295 new_nonzero_groups_list = \ 296 copy.copy(zero_blocks_list), \ 297 copy.copy(nonzero_blocks_list),\ 298 copy.copy(nonzero_groups_list) 299 return new_zero_blocks_list, new_nonzero_blocks_list, \ 300 new_nonzero_groups_list 301 302 @staticmethod 303 def get_file_data(*args): 304 """ 305 Get the file data. 306 block_size: blocksize, 307 chunk_start: the start position of chunk, 308 default_zero_block: default to zero blocks, 309 each_value: each_value, 310 file_pos: file position, 311 fill_data: data, 312 image_file_r: read file object, 313 :return data: Get the file data. 314 """ 315 block_size, chunk_start, default_zero_block, each_value, \ 316 file_pos, fill_data, image_file_r = args 317 if file_pos is not None: 318 file_pos += (each_value - chunk_start) * block_size 319 image_file_r.seek(file_pos, os.SEEK_SET) 320 data = image_file_r.read(block_size) 321 else: 322 if fill_data == default_zero_block[:4]: 323 data = default_zero_block 324 else: 325 data = None 326 return data 327 328 def range_sha256(self, ranges): 329 """ 330 range sha256 hash content 331 :param ranges: ranges value 332 :return: 333 """ 334 hash_obj = sha256() 335 for data in self.__get_blocks_set_data(ranges): 336 hash_obj.update(data) 337 return hash_obj.hexdigest() 338 339 def write_range_data_2_fd(self, ranges, file_obj): 340 """ 341 write range data to fd 342 :param ranges: ranges obj 343 :param file_obj: file obj 344 :return: 345 """ 346 for data in self.__get_blocks_set_data(ranges): 347 file_obj.write(data) 348 349 def get_ranges(self, ranges): 350 """ 351 get ranges value 352 :param ranges: ranges 353 :return: ranges value 354 """ 355 return [each_data for each_data in self.__get_blocks_set_data(ranges)] 356 357 def __get_blocks_set_data(self, blocks_set_data): 358 """ 359 Get the range data. 360 """ 361 with open(self.image_path, 'rb') as f_r: 362 for start, end in blocks_set_data: 363 diff_value = end - start 364 idx = bisect.bisect_right(self.offset_index, start) - 1 365 chunk_start, chunk_len, file_pos, fill_data = \ 366 self.offset_value_list[idx] 367 368 remain = chunk_len - (start - chunk_start) 369 this_read = min(remain, diff_value) 370 if file_pos is not None: 371 pos = file_pos + ((start - chunk_start) * self.block_size) 372 f_r.seek(pos, os.SEEK_SET) 373 yield f_r.read(this_read * self.block_size) 374 else: 375 yield fill_data * (this_read * (self.block_size >> 2)) 376 diff_value -= this_read 377 378 while diff_value > 0: 379 idx += 1 380 chunk_start, chunk_len, file_pos, fill_data = \ 381 self.offset_value_list[idx] 382 this_read = min(chunk_len, diff_value) 383 if file_pos is not None: 384 f_r.seek(file_pos, os.SEEK_SET) 385 yield f_r.read(this_read * self.block_size) 386 else: 387 yield fill_data * (this_read * (self.block_size >> 2)) 388 diff_value -= this_read