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. 16 17import multiprocessing 18import subprocess 19import tempfile 20import zipfile 21from ctypes import pointer 22from log_exception import UPDATE_LOGGER 23from blocks_manager import BlocksManager 24from transfers_manager import ActionType 25from update_package import PkgHeader 26from update_package import PkgComponent 27from utils import OPTIONS_MANAGER 28from utils import ON_SERVER 29from utils import DIFF_EXE_PATH 30 31NEW_DAT = "new.dat" 32PATCH_DAT = "patch.dat" 33TRANSFER_LIST = "transfer.list" 34 35 36class PatchProcess: 37 def __init__(self, partition, tgt_image, src_image, 38 actions_list): 39 self.actions_list = actions_list 40 self.worker_threads = multiprocessing.cpu_count() // 2 41 self.partition = partition 42 self.tgt_img_obj = tgt_image 43 self.src_img_obj = src_image 44 self.version = 1 45 self.touched_src_ranges = BlocksManager() 46 self.touched_src_sha256 = None 47 self.package_patch_zip = PackagePatchZip(partition) 48 49 def patch_process(self): 50 """ 51 Generate patches through calculation. 52 """ 53 UPDATE_LOGGER.print_log("Patch Process!") 54 55 new_dat_file_obj, patch_dat_file_obj, transfer_list_file_obj = \ 56 self.package_patch_zip.get_file_obj() 57 58 stashes = {} 59 total_blocks_count = 0 60 stashed_blocks = 0 61 max_stashed_blocks = 0 62 transfer_content = ["%d\n" % self.version, "TOTAL_MARK\n", 63 "0\n", "MAX_STASH_MARK\n"] 64 65 diff_offset = 0 66 for each_action in self.actions_list: 67 max_stashed_blocks, stashed_blocks = self.add_stash_command( 68 each_action, max_stashed_blocks, stashed_blocks, stashes, 69 transfer_content) 70 71 free_commands_list, free_size, src_str_list = \ 72 self.add_free_command(each_action, stashes) 73 74 src_str = " ".join(src_str_list) 75 tgt_size = each_action.tgt_block_set.size() 76 77 if each_action.type_str == ActionType.ZERO: 78 total_blocks_count = \ 79 self.apply_zero_type(each_action, total_blocks_count, 80 transfer_content) 81 elif each_action.type_str == ActionType.NEW: 82 total_blocks_count = \ 83 self.apply_new_type(each_action, new_dat_file_obj, 84 tgt_size, total_blocks_count, 85 transfer_content) 86 elif each_action.type_str == ActionType.DIFFERENT: 87 max_stashed_blocks, total_blocks_count, diff_offset = \ 88 self.apply_diff_style( 89 diff_offset, each_action, max_stashed_blocks, 90 patch_dat_file_obj, src_str, stashed_blocks, tgt_size, 91 total_blocks_count, transfer_content) 92 else: 93 UPDATE_LOGGER.print_log("Unknown action type: %s!" % 94 each_action.type_str) 95 raise RuntimeError 96 if free_commands_list: 97 transfer_content.append("".join(free_commands_list)) 98 stashed_blocks -= free_size 99 100 self.after_for_process(max_stashed_blocks, total_blocks_count, 101 transfer_content, transfer_list_file_obj) 102 103 def apply_new_type(self, each_action, new_dat_file_obj, tgt_size, 104 total_blocks_count, transfer_content): 105 self.tgt_img_obj.write_range_data_2_fd( 106 each_action.tgt_block_set, new_dat_file_obj) 107 UPDATE_LOGGER.print_log("%7s %s %s" % ( 108 each_action.type_str, each_action.tgt_name, 109 str(each_action.tgt_block_set))) 110 temp_size = self.write_split_transfers( 111 transfer_content, 112 each_action.type_str, each_action.tgt_block_set) 113 if tgt_size != temp_size: 114 raise RuntimeError 115 total_blocks_count += temp_size 116 return total_blocks_count 117 118 def apply_zero_type(self, each_action, total_blocks_count, 119 transfer_content): 120 UPDATE_LOGGER.print_log("%7s %s %s" % ( 121 each_action.type_str, each_action.tgt_name, 122 str(each_action.tgt_block_set))) 123 to_zero = \ 124 each_action.tgt_block_set.get_subtract_with_other( 125 each_action.src_block_set) 126 if self.write_split_transfers(transfer_content, each_action.type_str, 127 to_zero) != to_zero.size(): 128 raise RuntimeError 129 total_blocks_count += to_zero.size() 130 return total_blocks_count 131 132 def apply_diff_style(self, *args): 133 """ 134 Process actions of the diff type. 135 """ 136 diff_offset, each_action, max_stashed_blocks,\ 137 patch_dat_file_obj, src_str, stashed_blocks, tgt_size,\ 138 total_blocks_count, transfer_content = args 139 if self.tgt_img_obj. \ 140 range_sha256(each_action.tgt_block_set) == \ 141 self.src_img_obj.\ 142 range_sha256(each_action.src_block_set): 143 each_action.type_str = ActionType.MOVE 144 UPDATE_LOGGER.print_log("%7s %s %s (from %s %s)" % ( 145 each_action.type_str, each_action.tgt_name, 146 str(each_action.tgt_block_set), 147 each_action.src_name, 148 str(each_action.src_block_set))) 149 150 max_stashed_blocks, total_blocks_count = \ 151 self.add_move_command( 152 each_action, max_stashed_blocks, src_str, 153 stashed_blocks, tgt_size, total_blocks_count, 154 transfer_content) 155 else: 156 do_pkg_diff, patch_value = self.compute_diff_patch( 157 each_action, patch_dat_file_obj) 158 159 if each_action.src_block_set.is_overlaps( 160 each_action.tgt_block_set): 161 temp_stash_usage = \ 162 stashed_blocks + each_action.src_block_set.size() 163 if temp_stash_usage > max_stashed_blocks: 164 max_stashed_blocks = temp_stash_usage 165 166 self.add_diff_command(diff_offset, do_pkg_diff, 167 each_action, patch_value, src_str, 168 transfer_content) 169 170 diff_offset += len(patch_value) 171 total_blocks_count += tgt_size 172 return max_stashed_blocks, total_blocks_count, diff_offset 173 174 def after_for_process(self, max_stashed_blocks, total_blocks_count, 175 transfer_content, transfer_list_file_obj): 176 """ 177 Implement processing after cyclical actions_list processing. 178 :param max_stashed_blocks: maximum number of stashed blocks in actions 179 :param total_blocks_count: total number of blocks 180 :param transfer_content: transfer content 181 :param transfer_list_file_obj: transfer file object 182 :return: 183 """ 184 self.touched_src_sha256 = self.src_img_obj.range_sha256( 185 self.touched_src_ranges) 186 if self.tgt_img_obj.extended_range: 187 if self.write_split_transfers( 188 transfer_content, ActionType.ZERO, 189 self.tgt_img_obj.extended_range) != \ 190 self.tgt_img_obj.extended_range.size(): 191 raise RuntimeError 192 total_blocks_count += self.tgt_img_obj.extended_range.size() 193 all_tgt = BlocksManager( 194 range_data=(0, self.tgt_img_obj.total_blocks)) 195 all_tgt_minus_extended = all_tgt.get_subtract_with_other( 196 self.tgt_img_obj.extended_range) 197 new_not_care = all_tgt_minus_extended.get_subtract_with_other( 198 self.tgt_img_obj.care_block_range) 199 self.add_erase_content(new_not_care, transfer_content) 200 transfer_content = self.get_transfer_content( 201 max_stashed_blocks, total_blocks_count, transfer_content) 202 transfer_list_file_obj.write(transfer_content.encode()) 203 204 @staticmethod 205 def get_transfer_content(max_stashed_blocks, total_blocks_count, 206 transfer_content): 207 """ 208 Get the tranfer content. 209 """ 210 transfer_content = ''.join(transfer_content) 211 transfer_content = \ 212 transfer_content.replace("TOTAL_MARK", str(total_blocks_count)) 213 transfer_content = \ 214 transfer_content.replace("MAX_STASH_MARK", str(max_stashed_blocks)) 215 transfer_content = \ 216 transfer_content.replace("ActionType.MOVE", "move") 217 transfer_content = \ 218 transfer_content.replace("ActionType.ZERO", "zero") 219 transfer_content = \ 220 transfer_content.replace("ActionType.NEW", "new") 221 return transfer_content 222 223 def add_diff_command(self, *args): 224 """ 225 Add the diff command. 226 """ 227 diff_offset, do_pkg_diff, each_action,\ 228 patch_value, src_str, transfer_content = args 229 self.touched_src_ranges = self.touched_src_ranges.get_union_with_other( 230 each_action.src_block_set) 231 diff_type = "pkgdiff" if do_pkg_diff else "bsdiff" 232 transfer_content.append("%s %d %d %s %s %s %s\n" % ( 233 diff_type, 234 diff_offset, len(patch_value), 235 self.src_img_obj.range_sha256(each_action.src_block_set), 236 self.tgt_img_obj.range_sha256(each_action.tgt_block_set), 237 each_action.tgt_block_set.to_string_raw(), src_str)) 238 239 def compute_diff_patch(self, each_action, patch_dat_file_obj): 240 """ 241 Run the command to calculate the differential patch. 242 """ 243 src_file_obj = \ 244 tempfile.NamedTemporaryFile(prefix="src-", mode='wb') 245 self.src_img_obj.write_range_data_2_fd( 246 each_action.src_block_set, src_file_obj) 247 src_file_obj.seek(0) 248 tgt_file_obj = tempfile.NamedTemporaryFile( 249 prefix="tgt-", mode='wb') 250 self.tgt_img_obj.write_range_data_2_fd( 251 each_action.tgt_block_set, tgt_file_obj) 252 tgt_file_obj.seek(0) 253 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 254 src_file_obj) 255 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 256 tgt_file_obj) 257 do_pkg_diff = True 258 try: 259 patch_value, do_pkg_diff = self.apply_compute_patch( 260 src_file_obj.name, tgt_file_obj.name, do_pkg_diff) 261 src_file_obj.close() 262 tgt_file_obj.close() 263 except ValueError: 264 UPDATE_LOGGER.print_log("Patch process Failed!") 265 UPDATE_LOGGER.print_log("%7s %s %s (from %s %s)" % ( 266 each_action.type_str, each_action.tgt_name, 267 str(each_action.tgt_block_set), 268 each_action.src_name, 269 str(each_action.src_block_set)), 270 UPDATE_LOGGER.ERROR_LOG) 271 raise ValueError 272 patch_dat_file_obj.write(patch_value) 273 return do_pkg_diff, patch_value 274 275 def add_move_command(self, *args): 276 """ 277 Add the move command. 278 """ 279 each_action, max_stashed_blocks, src_str,\ 280 stashed_blocks, tgt_size, total_blocks_count,\ 281 transfer_content = args 282 src_block_set = each_action.src_block_set 283 tgt_block_set = each_action.tgt_block_set 284 if src_block_set != tgt_block_set: 285 if src_block_set.is_overlaps(tgt_block_set): 286 temp_stash_usage = stashed_blocks + \ 287 src_block_set.size() 288 if temp_stash_usage > max_stashed_blocks: 289 max_stashed_blocks = temp_stash_usage 290 291 self.touched_src_ranges = \ 292 self.touched_src_ranges.get_union_with_other(src_block_set) 293 294 transfer_content.append( 295 "{type_str} {tgt_hash} {tgt_string} {src_str}\n". 296 format(type_str=each_action.type_str, 297 tgt_hash=self.tgt_img_obj. 298 range_sha256(each_action.tgt_block_set), 299 tgt_string=tgt_block_set.to_string_raw(), 300 src_str=src_str)) 301 total_blocks_count += tgt_size 302 return max_stashed_blocks, total_blocks_count 303 304 def add_free_command(self, each_action, stashes): 305 """ 306 Add the free command. 307 :param each_action: action object to be processed 308 :param stashes: Stash dict 309 :return: free_commands_list, free_size, src_str_list 310 """ 311 free_commands_list = [] 312 free_size = 0 313 src_blocks_size = each_action.src_block_set.size() 314 src_str_list = [str(src_blocks_size)] 315 un_stashed_src_ranges = each_action.src_block_set 316 mapped_stashes = [] 317 for _, each_stash_before in each_action.use_stash: 318 un_stashed_src_ranges = \ 319 un_stashed_src_ranges.get_subtract_with_other( 320 each_stash_before) 321 src_range_sha = \ 322 self.src_img_obj.range_sha256(each_stash_before) 323 each_stash_before = \ 324 each_action.src_block_set.get_map_within(each_stash_before) 325 mapped_stashes.append(each_stash_before) 326 if src_range_sha not in stashes: 327 raise RuntimeError 328 src_str_list.append( 329 "%s:%s" % (src_range_sha, each_stash_before.to_string_raw())) 330 stashes[src_range_sha] -= 1 331 if stashes[src_range_sha] == 0: 332 free_commands_list.append("free %s\n" % (src_range_sha,)) 333 free_size += each_stash_before.size() 334 stashes.pop(src_range_sha) 335 self.apply_stashed_range(each_action, mapped_stashes, src_blocks_size, 336 src_str_list, un_stashed_src_ranges) 337 return free_commands_list, free_size, src_str_list 338 339 def apply_stashed_range(self, *args): 340 each_action, mapped_stashes, src_blocks_size,\ 341 src_str_list, un_stashed_src_ranges = args 342 if un_stashed_src_ranges.size() != 0: 343 src_str_list.insert(1, un_stashed_src_ranges.to_string_raw()) 344 if each_action.use_stash: 345 mapped_un_stashed = each_action.src_block_set.get_map_within( 346 un_stashed_src_ranges) 347 src_str_list.insert(2, mapped_un_stashed.to_string_raw()) 348 mapped_stashes.append(mapped_un_stashed) 349 self.check_partition( 350 BlocksManager(range_data=(0, src_blocks_size)), 351 mapped_stashes) 352 else: 353 src_str_list.insert(1, "-") 354 self.check_partition( 355 BlocksManager(range_data=(0, src_blocks_size)), mapped_stashes) 356 357 def add_stash_command(self, each_action, max_stashed_blocks, 358 stashed_blocks, stashes, transfer_content): 359 """ 360 Add the stash command. 361 :param each_action: action object to be processed 362 :param max_stashed_blocks: number of max stash blocks in all actions 363 :param stashed_blocks: number of stash blocks 364 :param stashes: Stash dict 365 :param transfer_content: transfer content list 366 :return: max_stashed_blocks, stashed_blocks 367 """ 368 for _, each_stash_before in each_action.stash_before: 369 src_range_sha = \ 370 self.src_img_obj.range_sha256(each_stash_before) 371 if src_range_sha in stashes: 372 stashes[src_range_sha] += 1 373 else: 374 stashes[src_range_sha] = 1 375 stashed_blocks += each_stash_before.size() 376 self.touched_src_ranges = \ 377 self.touched_src_ranges.\ 378 get_union_with_other(each_stash_before) 379 transfer_content.append("stash %s %s\n" % ( 380 src_range_sha, each_stash_before.to_string_raw())) 381 if stashed_blocks > max_stashed_blocks: 382 max_stashed_blocks = stashed_blocks 383 return max_stashed_blocks, stashed_blocks 384 385 def write_script(self, partition, script_check_cmd_list, 386 script_write_cmd_list, verse_script): 387 """ 388 Add command content to the script. 389 :param partition: image name 390 :param script_check_cmd_list: incremental check command list 391 :param script_write_cmd_list: incremental write command list 392 :param verse_script: verse script object 393 :return: 394 """ 395 ranges_str = self.touched_src_ranges.to_string_raw() 396 expected_sha = self.touched_src_sha256 397 398 sha_check_cmd = verse_script.sha_check( 399 ranges_str, expected_sha, partition) 400 401 first_block_check_cmd = verse_script.first_block_check(partition) 402 403 abort_cmd = verse_script.abort(partition) 404 405 cmd = 'if ({sha_check_cmd} != 0 || ' \ 406 '{first_block_check_cmd} != 0)' \ 407 '{{\n {abort_cmd}}}\n'.format( 408 sha_check_cmd=sha_check_cmd, 409 first_block_check_cmd=first_block_check_cmd, 410 abort_cmd=abort_cmd) 411 412 script_check_cmd_list.append(cmd) 413 414 block_update_cmd = verse_script.block_update(partition) 415 416 cmd = '%s_WRITE_FLAG%s' % (partition, block_update_cmd) 417 script_write_cmd_list.append(cmd) 418 419 def add_erase_content(self, new_not_care, transfer_content): 420 """ 421 Add the erase command. 422 :param new_not_care: blocks that don't need to be cared about 423 :param transfer_content: transfer content list 424 :return: 425 """ 426 erase_first = new_not_care.\ 427 get_subtract_with_other(self.touched_src_ranges) 428 if erase_first.size() != 0: 429 transfer_content.insert( 430 4, "erase %s\n" % (erase_first.to_string_raw(),)) 431 erase_last = new_not_care.get_subtract_with_other(erase_first) 432 if erase_last.size() != 0: 433 transfer_content.append( 434 "erase %s\n" % (erase_last.to_string_raw(),)) 435 436 @staticmethod 437 def check_partition(total, seq): 438 so_far = BlocksManager() 439 for i in seq: 440 if so_far.is_overlaps(i): 441 raise RuntimeError 442 so_far = so_far.get_union_with_other(i) 443 if so_far != total: 444 raise RuntimeError 445 446 @staticmethod 447 def write_split_transfers(transfer_content, type_str, target_blocks): 448 """ 449 Limit the size of operand in command 'new' and 'zero' to 1024 blocks. 450 :param transfer_content: transfer content list 451 :param type_str: type of the action to be processed. 452 :param target_blocks: BlocksManager of the target blocks 453 :return: total 454 """ 455 if type_str not in (ActionType.NEW, ActionType.ZERO): 456 raise RuntimeError 457 blocks_limit = 1024 458 total = 0 459 while target_blocks.size() != 0: 460 blocks_to_write = target_blocks.get_first_block_obj(blocks_limit) 461 transfer_content.append( 462 "%s %s\n" % (type_str, blocks_to_write.to_string_raw())) 463 total += blocks_to_write.size() 464 target_blocks = \ 465 target_blocks.get_subtract_with_other(blocks_to_write) 466 return total 467 468 @staticmethod 469 def apply_compute_patch(src_file, tgt_file, pkgdiff=False): 470 """ 471 Add command content to the script. 472 :param src_file: source file name 473 :param tgt_file: target file name 474 :param pkgdiff: whether to execute pkgdiff judgment 475 :return: 476 """ 477 patch_file_obj = \ 478 tempfile.NamedTemporaryFile(prefix="patch-", mode='wb') 479 480 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 481 patch_file_obj) 482 cmd = [DIFF_EXE_PATH] if pkgdiff else [DIFF_EXE_PATH, '-b', '1'] 483 484 cmd.extend(['-s', src_file, '-d', tgt_file, 485 '-p', patch_file_obj.name, '-l', '4096']) 486 sub_p = subprocess.Popen(cmd, stdout=subprocess.PIPE, 487 stderr=subprocess.STDOUT) 488 output, _ = sub_p.communicate() 489 sub_p.wait() 490 patch_file_obj.seek(0) 491 492 if sub_p.returncode != 0: 493 raise ValueError(output) 494 495 with open(patch_file_obj.name, 'rb') as file_read: 496 patch_content = file_read.read() 497 return patch_content, pkgdiff 498 499 500class PackagePatchZip: 501 """ 502 Compress the patch file generated by the 503 differential calculation as *.zip file. 504 """ 505 def __init__(self, partition): 506 self.partition = partition 507 self.partition_new_dat_file_name = "%s.%s" % (partition, NEW_DAT) 508 self.partition_patch_dat_file_name = "%s.%s" % (partition, PATCH_DAT) 509 self.partition_transfer_file_name = "%s.%s" % (partition, TRANSFER_LIST) 510 511 self.new_dat_file_obj = tempfile.NamedTemporaryFile( 512 dir=OPTIONS_MANAGER.target_package, prefix="%s-" % NEW_DAT, mode='wb') 513 self.patch_dat_file_obj = tempfile.NamedTemporaryFile( 514 dir=OPTIONS_MANAGER.target_package, prefix="%s-" % PATCH_DAT, mode='wb') 515 self.transfer_list_file_obj = tempfile.NamedTemporaryFile( 516 dir=OPTIONS_MANAGER.target_package, prefix="%s-" % TRANSFER_LIST, mode='wb') 517 518 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 519 self.new_dat_file_obj) 520 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 521 self.patch_dat_file_obj) 522 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 523 self.transfer_list_file_obj) 524 525 self.partition_file_obj = tempfile.NamedTemporaryFile( 526 dir=OPTIONS_MANAGER.target_package, prefix="partition_patch-") 527 OPTIONS_MANAGER.incremental_image_file_obj_list.append( 528 self.partition_file_obj) 529 530 def get_file_obj(self): 531 """ 532 Obtain file objects. 533 """ 534 return self.new_dat_file_obj, self.patch_dat_file_obj, \ 535 self.transfer_list_file_obj 536 537 def package_patch_zip(self): 538 """ 539 Compress the partition diff patch calculation data as *.zip package. 540 """ 541 self.partition_file_obj.seek(0) 542 self.patch_dat_file_obj.seek(0) 543 self.new_dat_file_obj.seek(0) 544 self.transfer_list_file_obj.seek(0) 545 546 with zipfile.ZipFile(self.partition_file_obj.name, 'w', zipfile.ZIP_DEFLATED) as zip_file: 547 # add new.dat to {partition}.zip 548 zip_file.write(self.new_dat_file_obj.name, self.partition_new_dat_file_name) 549 # add patch.dat to {partition}.zip 550 zip_file.write(self.patch_dat_file_obj.name, self.partition_patch_dat_file_name) 551 # add transfer.list to {partition}.zip 552 zip_file.write(self.transfer_list_file_obj.name, self.partition_transfer_file_name) 553 UPDATE_LOGGER.print_log("Create %s.zip success!" % self.partition) 554