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