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, 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, 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 elif each_action .tgt_block_set.size() > 125 * 1024: # target_file_size > 125 * 1024 * 4KB = 500M 156 each_action.type_str = ActionType.NEW 157 new_dat_file_obj, patch_dat_file_obj, transfer_list_file_obj = \ 158 self.package_patch_zip.get_file_obj() 159 total_blocks_count = \ 160 self.apply_new_type(each_action, new_dat_file_obj, 161 tgt_size, total_blocks_count, 162 transfer_content) 163 else: 164 do_pkg_diff, patch_value = self.compute_diff_patch( 165 each_action, patch_dat_file_obj) 166 167 if each_action.src_block_set.is_overlaps( 168 each_action.tgt_block_set): 169 stashed_blocks = \ 170 stashed_blocks + each_action.src_block_set.size() 171 if stashed_blocks > max_stashed_blocks: 172 max_stashed_blocks = stashed_blocks 173 174 self.add_diff_command(diff_offset, do_pkg_diff, 175 each_action, patch_value, src_str, 176 transfer_content) 177 178 diff_offset += len(patch_value) 179 total_blocks_count += tgt_size 180 return max_stashed_blocks, stashed_blocks, total_blocks_count, diff_offset 181 182 def after_for_process(self, max_stashed_blocks, total_blocks_count, 183 transfer_content, transfer_list_file_obj): 184 """ 185 Implement processing after cyclical actions_list processing. 186 :param max_stashed_blocks: maximum number of stashed blocks in actions 187 :param total_blocks_count: total number of blocks 188 :param transfer_content: transfer content 189 :param transfer_list_file_obj: transfer file object 190 :return: 191 """ 192 self.touched_src_sha256 = self.src_img_obj.range_sha256( 193 self.touched_src_ranges) 194 if self.tgt_img_obj.extended_range: 195 if self.write_split_transfers( 196 transfer_content, ActionType.ZERO, 197 self.tgt_img_obj.extended_range) != \ 198 self.tgt_img_obj.extended_range.size(): 199 raise RuntimeError 200 total_blocks_count += self.tgt_img_obj.extended_range.size() 201 all_tgt = BlocksManager( 202 range_data=(0, self.tgt_img_obj.total_blocks)) 203 all_tgt_minus_extended = all_tgt.get_subtract_with_other( 204 self.tgt_img_obj.extended_range) 205 new_not_care = all_tgt_minus_extended.get_subtract_with_other( 206 self.tgt_img_obj.care_block_range) 207 self.add_erase_content(new_not_care, transfer_content) 208 transfer_content = self.get_transfer_content( 209 max_stashed_blocks, total_blocks_count, transfer_content) 210 transfer_list_file_obj.write(transfer_content.encode()) 211 OPTIONS_MANAGER.max_stash_size = max(max_stashed_blocks * 4096, OPTIONS_MANAGER.max_stash_size) 212 213 @staticmethod 214 def get_transfer_content(max_stashed_blocks, total_blocks_count, 215 transfer_content): 216 """ 217 Get the tranfer content. 218 """ 219 transfer_content = ''.join(transfer_content) 220 transfer_content = \ 221 transfer_content.replace("TOTAL_MARK", str(total_blocks_count)) 222 transfer_content = \ 223 transfer_content.replace("MAX_STASH_MARK", str(max_stashed_blocks)) 224 transfer_content = \ 225 transfer_content.replace("ActionType.MOVE", "move") 226 transfer_content = \ 227 transfer_content.replace("ActionType.ZERO", "zero") 228 transfer_content = \ 229 transfer_content.replace("ActionType.NEW", "new") 230 return transfer_content 231 232 def add_diff_command(self, *args): 233 """ 234 Add the diff command. 235 """ 236 diff_offset, do_pkg_diff, each_action,\ 237 patch_value, src_str, transfer_content = args 238 self.touched_src_ranges = self.touched_src_ranges.get_union_with_other( 239 each_action.src_block_set) 240 diff_type = "pkgdiff" if do_pkg_diff else "bsdiff" 241 transfer_content.append("%s %d %d %s %s %s %s\n" % ( 242 diff_type, 243 diff_offset, len(patch_value), 244 self.src_img_obj.range_sha256(each_action.src_block_set), 245 self.tgt_img_obj.range_sha256(each_action.tgt_block_set), 246 each_action.tgt_block_set.to_string_raw(), src_str)) 247 248 def compute_diff_patch(self, each_action, patch_dat_file_obj): 249 """ 250 Run the command to calculate the differential patch. 251 """ 252 src_file_obj = \ 253 tempfile.NamedTemporaryFile(prefix="src-", mode='wb') 254 self.src_img_obj.write_range_data_2_fd( 255 each_action.src_block_set, src_file_obj) 256 src_file_obj.seek(0) 257 tgt_file_obj = tempfile.NamedTemporaryFile( 258 prefix="tgt-", mode='wb') 259 self.tgt_img_obj.write_range_data_2_fd( 260 each_action.tgt_block_set, tgt_file_obj) 261 tgt_file_obj.seek(0) 262 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 263 src_file_obj) 264 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 265 tgt_file_obj) 266 do_pkg_diff = True 267 try: 268 patch_value, do_pkg_diff = self.apply_compute_patch( 269 src_file_obj.name, tgt_file_obj.name, do_pkg_diff) 270 src_file_obj.close() 271 tgt_file_obj.close() 272 except ValueError: 273 UPDATE_LOGGER.print_log("Patch process Failed!") 274 UPDATE_LOGGER.print_log("%7s %s %s (from %s %s)" % ( 275 each_action.type_str, each_action.tgt_name, 276 str(each_action.tgt_block_set), 277 each_action.src_name, 278 str(each_action.src_block_set)), 279 UPDATE_LOGGER.ERROR_LOG) 280 raise ValueError 281 patch_dat_file_obj.write(patch_value) 282 return do_pkg_diff, patch_value 283 284 def add_move_command(self, *args): 285 """ 286 Add the move command. 287 """ 288 each_action, max_stashed_blocks, src_str,\ 289 stashed_blocks, tgt_size, total_blocks_count,\ 290 transfer_content = args 291 src_block_set = each_action.src_block_set 292 tgt_block_set = each_action.tgt_block_set 293 if src_block_set != tgt_block_set: 294 if src_block_set.is_overlaps(tgt_block_set): 295 stashed_blocks = stashed_blocks + \ 296 src_block_set.size() 297 if stashed_blocks > max_stashed_blocks: 298 max_stashed_blocks = stashed_blocks 299 300 self.touched_src_ranges = \ 301 self.touched_src_ranges.get_union_with_other(src_block_set) 302 303 transfer_content.append( 304 "{type_str} {tgt_hash} {tgt_string} {src_str}\n". 305 format(type_str=each_action.type_str, 306 tgt_hash=self.tgt_img_obj. 307 range_sha256(each_action.tgt_block_set), 308 tgt_string=tgt_block_set.to_string_raw(), 309 src_str=src_str)) 310 total_blocks_count += tgt_size 311 return max_stashed_blocks, stashed_blocks, total_blocks_count 312 313 def add_free_command(self, each_action, stashes): 314 """ 315 Add the free command. 316 :param each_action: action object to be processed 317 :param stashes: Stash dict 318 :return: free_commands_list, free_size, src_str_list 319 """ 320 free_commands_list = [] 321 free_size = 0 322 src_blocks_size = each_action.src_block_set.size() 323 src_str_list = [str(src_blocks_size)] 324 un_stashed_src_ranges = each_action.src_block_set 325 mapped_stashes = [] 326 for _, each_stash_before in each_action.use_stash: 327 un_stashed_src_ranges = \ 328 un_stashed_src_ranges.get_subtract_with_other( 329 each_stash_before) 330 src_range_sha = \ 331 self.src_img_obj.range_sha256(each_stash_before) 332 each_stash_before = \ 333 each_action.src_block_set.get_map_within(each_stash_before) 334 mapped_stashes.append(each_stash_before) 335 if src_range_sha not in stashes: 336 raise RuntimeError 337 src_str_list.append( 338 "%s:%s" % (src_range_sha, each_stash_before.to_string_raw())) 339 stashes[src_range_sha] -= 1 340 if stashes[src_range_sha] == 0: 341 free_commands_list.append("free %s\n" % (src_range_sha,)) 342 free_size += each_stash_before.size() 343 stashes.pop(src_range_sha) 344 self.apply_stashed_range(each_action, mapped_stashes, src_blocks_size, 345 src_str_list, un_stashed_src_ranges) 346 return free_commands_list, free_size, src_str_list 347 348 def apply_stashed_range(self, *args): 349 each_action, mapped_stashes, src_blocks_size,\ 350 src_str_list, un_stashed_src_ranges = args 351 if un_stashed_src_ranges.size() != 0: 352 src_str_list.insert(1, un_stashed_src_ranges.to_string_raw()) 353 if each_action.use_stash: 354 mapped_un_stashed = each_action.src_block_set.get_map_within( 355 un_stashed_src_ranges) 356 src_str_list.insert(2, mapped_un_stashed.to_string_raw()) 357 mapped_stashes.append(mapped_un_stashed) 358 self.check_partition( 359 BlocksManager(range_data=(0, src_blocks_size)), 360 mapped_stashes) 361 else: 362 src_str_list.insert(1, "-") 363 self.check_partition( 364 BlocksManager(range_data=(0, src_blocks_size)), mapped_stashes) 365 366 def add_stash_command(self, each_action, max_stashed_blocks, 367 stashed_blocks, stashes, transfer_content): 368 """ 369 Add the stash command. 370 :param each_action: action object to be processed 371 :param max_stashed_blocks: number of max stash blocks in all actions 372 :param stashed_blocks: number of stash blocks 373 :param stashes: Stash dict 374 :param transfer_content: transfer content list 375 :return: max_stashed_blocks, stashed_blocks 376 """ 377 for _, each_stash_before in each_action.stash_before: 378 src_range_sha = \ 379 self.src_img_obj.range_sha256(each_stash_before) 380 if src_range_sha in stashes: 381 stashes[src_range_sha] += 1 382 else: 383 stashes[src_range_sha] = 1 384 stashed_blocks += each_stash_before.size() 385 self.touched_src_ranges = \ 386 self.touched_src_ranges.\ 387 get_union_with_other(each_stash_before) 388 transfer_content.append("stash %s %s\n" % ( 389 src_range_sha, each_stash_before.to_string_raw())) 390 if stashed_blocks > max_stashed_blocks: 391 max_stashed_blocks = stashed_blocks 392 return max_stashed_blocks, stashed_blocks 393 394 def write_script(self, partition, script_check_cmd_list, 395 script_write_cmd_list, verse_script): 396 """ 397 Add command content to the script. 398 :param partition: image name 399 :param script_check_cmd_list: incremental check command list 400 :param script_write_cmd_list: incremental write command list 401 :param verse_script: verse script object 402 :return: 403 """ 404 ranges_str = self.touched_src_ranges.to_string_raw() 405 expected_sha = self.touched_src_sha256 406 407 sha_check_cmd = verse_script.sha_check( 408 ranges_str, expected_sha, partition) 409 410 first_block_check_cmd = verse_script.first_block_check(partition) 411 412 abort_cmd = verse_script.abort(partition) 413 414 cmd = 'if ({sha_check_cmd} != 0 || ' \ 415 '{first_block_check_cmd} != 0)' \ 416 '{{\n {abort_cmd}}}\n'.format( 417 sha_check_cmd=sha_check_cmd, 418 first_block_check_cmd=first_block_check_cmd, 419 abort_cmd=abort_cmd) 420 421 script_check_cmd_list.append(cmd) 422 423 block_update_cmd = verse_script.block_update(partition) 424 425 cmd = '%s_WRITE_FLAG%s' % (partition, block_update_cmd) 426 script_write_cmd_list.append(cmd) 427 428 def add_erase_content(self, new_not_care, transfer_content): 429 """ 430 Add the erase command. 431 :param new_not_care: blocks that don't need to be cared about 432 :param transfer_content: transfer content list 433 :return: 434 """ 435 erase_first = new_not_care.\ 436 get_subtract_with_other(self.touched_src_ranges) 437 if erase_first.size() != 0: 438 transfer_content.insert( 439 4, "erase %s\n" % (erase_first.to_string_raw(),)) 440 erase_last = new_not_care.get_subtract_with_other(erase_first) 441 if erase_last.size() != 0: 442 transfer_content.append( 443 "erase %s\n" % (erase_last.to_string_raw(),)) 444 445 @staticmethod 446 def check_partition(total, seq): 447 so_far = BlocksManager() 448 for i in seq: 449 if so_far.is_overlaps(i): 450 raise RuntimeError 451 so_far = so_far.get_union_with_other(i) 452 if so_far != total: 453 raise RuntimeError 454 455 @staticmethod 456 def write_split_transfers(transfer_content, type_str, target_blocks): 457 """ 458 Limit the size of operand in command 'new' and 'zero' to 1024 blocks. 459 :param transfer_content: transfer content list 460 :param type_str: type of the action to be processed. 461 :param target_blocks: BlocksManager of the target blocks 462 :return: total 463 """ 464 if type_str not in (ActionType.NEW, ActionType.ZERO): 465 raise RuntimeError 466 blocks_limit = 1024 467 total = 0 468 while target_blocks.size() != 0: 469 blocks_to_write = target_blocks.get_first_block_obj(blocks_limit) 470 transfer_content.append( 471 "%s %s\n" % (type_str, blocks_to_write.to_string_raw())) 472 total += blocks_to_write.size() 473 target_blocks = \ 474 target_blocks.get_subtract_with_other(blocks_to_write) 475 return total 476 477 @staticmethod 478 def apply_compute_patch(src_file, tgt_file, pkgdiff=False): 479 """ 480 Add command content to the script. 481 :param src_file: source file name 482 :param tgt_file: target file name 483 :param pkgdiff: whether to execute pkgdiff judgment 484 :return: 485 """ 486 patch_file_obj = \ 487 tempfile.NamedTemporaryFile(prefix="patch-", mode='wb') 488 489 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 490 patch_file_obj) 491 cmd = [DIFF_EXE_PATH] if pkgdiff else [DIFF_EXE_PATH, '-b', '1'] 492 493 cmd.extend(['-s', src_file, '-d', tgt_file, 494 '-p', patch_file_obj.name, '-l', '4096']) 495 sub_p = subprocess.Popen(cmd, stdout=subprocess.PIPE, 496 stderr=subprocess.STDOUT) 497 output, _ = sub_p.communicate() 498 sub_p.wait() 499 patch_file_obj.seek(0) 500 501 if sub_p.returncode != 0: 502 raise ValueError(output) 503 504 with open(patch_file_obj.name, 'rb') as file_read: 505 patch_content = file_read.read() 506 return patch_content, pkgdiff 507 508 509class PackagePatchZip: 510 """ 511 Compress the patch file generated by the 512 differential calculation as *.zip file. 513 """ 514 def __init__(self, partition): 515 self.partition = partition 516 self.partition_new_dat_file_name = "%s.%s" % (partition, NEW_DAT) 517 self.partition_patch_dat_file_name = "%s.%s" % (partition, PATCH_DAT) 518 self.partition_transfer_file_name = "%s.%s" % (partition, TRANSFER_LIST) 519 520 self.new_dat_file_obj = tempfile.NamedTemporaryFile( 521 dir=OPTIONS_MANAGER.target_package, prefix="%s-" % NEW_DAT, mode='wb') 522 self.patch_dat_file_obj = tempfile.NamedTemporaryFile( 523 dir=OPTIONS_MANAGER.target_package, prefix="%s-" % PATCH_DAT, mode='wb') 524 self.transfer_list_file_obj = tempfile.NamedTemporaryFile( 525 dir=OPTIONS_MANAGER.target_package, prefix="%s-" % TRANSFER_LIST, mode='wb') 526 527 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 528 self.new_dat_file_obj) 529 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 530 self.patch_dat_file_obj) 531 OPTIONS_MANAGER.incremental_temp_file_obj_list.append( 532 self.transfer_list_file_obj) 533 534 self.partition_file_obj = tempfile.NamedTemporaryFile( 535 dir=OPTIONS_MANAGER.target_package, prefix="partition_patch-") 536 537 def get_file_obj(self): 538 """ 539 Obtain file objects. 540 """ 541 self.new_dat_file_obj.flush() 542 self.patch_dat_file_obj.flush() 543 self.transfer_list_file_obj.flush() 544 return self.new_dat_file_obj, self.patch_dat_file_obj, \ 545 self.transfer_list_file_obj 546 547 def package_block_patch(self, zip_file): 548 self.new_dat_file_obj.flush() 549 self.patch_dat_file_obj.flush() 550 self.transfer_list_file_obj.flush() 551 # add new.dat to ota.zip 552 zip_file.write(self.new_dat_file_obj.name, self.partition_new_dat_file_name) 553 # add patch.dat to ota.zip 554 zip_file.write(self.patch_dat_file_obj.name, self.partition_patch_dat_file_name) 555 # add transfer.list to ota.zip 556 zip_file.write(self.transfer_list_file_obj.name, self.partition_transfer_file_name) 557