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