common.py revision 24cd28032890b1ca7fbc34320a22bd575f59d641
1# Copyright (C) 2008 The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15import copy 16import errno 17import getopt 18import getpass 19import imp 20import os 21import platform 22import re 23import shutil 24import subprocess 25import sys 26import tempfile 27import threading 28import time 29import zipfile 30 31try: 32 from hashlib import sha1 as sha1 33except ImportError: 34 from sha import sha as sha1 35 36# missing in Python 2.4 and before 37if not hasattr(os, "SEEK_SET"): 38 os.SEEK_SET = 0 39 40class Options(object): pass 41OPTIONS = Options() 42OPTIONS.search_path = "out/host/linux-x86" 43OPTIONS.verbose = False 44OPTIONS.tempfiles = [] 45OPTIONS.device_specific = None 46OPTIONS.extras = {} 47OPTIONS.info_dict = None 48 49 50# Values for "certificate" in apkcerts that mean special things. 51SPECIAL_CERT_STRINGS = ("PRESIGNED", "EXTERNAL") 52 53 54class ExternalError(RuntimeError): pass 55 56 57def Run(args, **kwargs): 58 """Create and return a subprocess.Popen object, printing the command 59 line on the terminal if -v was specified.""" 60 if OPTIONS.verbose: 61 print " running: ", " ".join(args) 62 return subprocess.Popen(args, **kwargs) 63 64 65def CloseInheritedPipes(): 66 """ Gmake in MAC OS has file descriptor (PIPE) leak. We close those fds 67 before doing other work.""" 68 if platform.system() != "Darwin": 69 return 70 for d in range(3, 1025): 71 try: 72 stat = os.fstat(d) 73 if stat is not None: 74 pipebit = stat[0] & 0x1000 75 if pipebit != 0: 76 os.close(d) 77 except OSError: 78 pass 79 80 81def LoadInfoDict(zip): 82 """Read and parse the META/misc_info.txt key/value pairs from the 83 input target files and return a dict.""" 84 85 d = {} 86 try: 87 for line in zip.read("META/misc_info.txt").split("\n"): 88 line = line.strip() 89 if not line or line.startswith("#"): continue 90 k, v = line.split("=", 1) 91 d[k] = v 92 except KeyError: 93 # ok if misc_info.txt doesn't exist 94 pass 95 96 # backwards compatibility: These values used to be in their own 97 # files. Look for them, in case we're processing an old 98 # target_files zip. 99 100 if "mkyaffs2_extra_flags" not in d: 101 try: 102 d["mkyaffs2_extra_flags"] = zip.read("META/mkyaffs2-extra-flags.txt").strip() 103 except KeyError: 104 # ok if flags don't exist 105 pass 106 107 if "recovery_api_version" not in d: 108 try: 109 d["recovery_api_version"] = zip.read("META/recovery-api-version.txt").strip() 110 except KeyError: 111 raise ValueError("can't find recovery API version in input target-files") 112 113 if "tool_extensions" not in d: 114 try: 115 d["tool_extensions"] = zip.read("META/tool-extensions.txt").strip() 116 except KeyError: 117 # ok if extensions don't exist 118 pass 119 120 try: 121 data = zip.read("META/imagesizes.txt") 122 for line in data.split("\n"): 123 if not line: continue 124 name, value = line.split(" ", 1) 125 if not value: continue 126 if name == "blocksize": 127 d[name] = value 128 else: 129 d[name + "_size"] = value 130 except KeyError: 131 pass 132 133 def makeint(key): 134 if key in d: 135 d[key] = int(d[key], 0) 136 137 makeint("recovery_api_version") 138 makeint("blocksize") 139 makeint("system_size") 140 makeint("userdata_size") 141 makeint("cache_size") 142 makeint("recovery_size") 143 makeint("boot_size") 144 145 d["fstab"] = LoadRecoveryFSTab(zip) 146 return d 147 148def LoadRecoveryFSTab(zip): 149 class Partition(object): 150 pass 151 152 try: 153 data = zip.read("RECOVERY/RAMDISK/etc/recovery.fstab") 154 except KeyError: 155 print "Warning: could not find RECOVERY/RAMDISK/etc/recovery.fstab in %s." % zip 156 data = "" 157 158 d = {} 159 for line in data.split("\n"): 160 line = line.strip() 161 if not line or line.startswith("#"): continue 162 pieces = line.split() 163 if not (3 <= len(pieces) <= 4): 164 raise ValueError("malformed recovery.fstab line: \"%s\"" % (line,)) 165 166 p = Partition() 167 p.mount_point = pieces[0] 168 p.fs_type = pieces[1] 169 p.device = pieces[2] 170 p.length = 0 171 options = None 172 if len(pieces) >= 4: 173 if pieces[3].startswith("/"): 174 p.device2 = pieces[3] 175 if len(pieces) >= 5: 176 options = pieces[4] 177 else: 178 p.device2 = None 179 options = pieces[3] 180 else: 181 p.device2 = None 182 183 if options: 184 options = options.split(",") 185 for i in options: 186 if i.startswith("length="): 187 p.length = int(i[7:]) 188 else: 189 print "%s: unknown option \"%s\"" % (p.mount_point, i) 190 191 d[p.mount_point] = p 192 return d 193 194 195def DumpInfoDict(d): 196 for k, v in sorted(d.items()): 197 print "%-25s = (%s) %s" % (k, type(v).__name__, v) 198 199def BuildBootableImage(sourcedir, fs_config_file, info_dict=None): 200 """Take a kernel, cmdline, and ramdisk directory from the input (in 201 'sourcedir'), and turn them into a boot image. Return the image 202 data, or None if sourcedir does not appear to contains files for 203 building the requested image.""" 204 205 if (not os.access(os.path.join(sourcedir, "RAMDISK"), os.F_OK) or 206 not os.access(os.path.join(sourcedir, "kernel"), os.F_OK)): 207 return None 208 209 if info_dict is None: 210 info_dict = OPTIONS.info_dict 211 212 ramdisk_img = tempfile.NamedTemporaryFile() 213 img = tempfile.NamedTemporaryFile() 214 215 if os.access(fs_config_file, os.F_OK): 216 cmd = ["mkbootfs", "-f", fs_config_file, os.path.join(sourcedir, "RAMDISK")] 217 else: 218 cmd = ["mkbootfs", os.path.join(sourcedir, "RAMDISK")] 219 p1 = Run(cmd, stdout=subprocess.PIPE) 220 p2 = Run(["minigzip"], 221 stdin=p1.stdout, stdout=ramdisk_img.file.fileno()) 222 223 p2.wait() 224 p1.wait() 225 assert p1.returncode == 0, "mkbootfs of %s ramdisk failed" % (targetname,) 226 assert p2.returncode == 0, "minigzip of %s ramdisk failed" % (targetname,) 227 228 cmd = ["mkbootimg", "--kernel", os.path.join(sourcedir, "kernel")] 229 230 fn = os.path.join(sourcedir, "cmdline") 231 if os.access(fn, os.F_OK): 232 cmd.append("--cmdline") 233 cmd.append(open(fn).read().rstrip("\n")) 234 235 fn = os.path.join(sourcedir, "base") 236 if os.access(fn, os.F_OK): 237 cmd.append("--base") 238 cmd.append(open(fn).read().rstrip("\n")) 239 240 fn = os.path.join(sourcedir, "pagesize") 241 if os.access(fn, os.F_OK): 242 cmd.append("--pagesize") 243 cmd.append(open(fn).read().rstrip("\n")) 244 245 args = info_dict.get("mkbootimg_args", None) 246 if args and args.strip(): 247 cmd.extend(args.split()) 248 249 cmd.extend(["--ramdisk", ramdisk_img.name, 250 "--output", img.name]) 251 252 p = Run(cmd, stdout=subprocess.PIPE) 253 p.communicate() 254 assert p.returncode == 0, "mkbootimg of %s image failed" % ( 255 os.path.basename(sourcedir),) 256 257 img.seek(os.SEEK_SET, 0) 258 data = img.read() 259 260 ramdisk_img.close() 261 img.close() 262 263 return data 264 265 266def GetBootableImage(name, prebuilt_name, unpack_dir, tree_subdir, 267 info_dict=None): 268 """Return a File object (with name 'name') with the desired bootable 269 image. Look for it in 'unpack_dir'/BOOTABLE_IMAGES under the name 270 'prebuilt_name', otherwise construct it from the source files in 271 'unpack_dir'/'tree_subdir'.""" 272 273 prebuilt_path = os.path.join(unpack_dir, "BOOTABLE_IMAGES", prebuilt_name) 274 if os.path.exists(prebuilt_path): 275 print "using prebuilt %s..." % (prebuilt_name,) 276 return File.FromLocalFile(name, prebuilt_path) 277 else: 278 print "building image from target_files %s..." % (tree_subdir,) 279 fs_config = "META/" + tree_subdir.lower() + "_filesystem_config.txt" 280 return File(name, BuildBootableImage(os.path.join(unpack_dir, tree_subdir), 281 os.path.join(unpack_dir, fs_config), 282 info_dict)) 283 284 285def UnzipTemp(filename, pattern=None): 286 """Unzip the given archive into a temporary directory and return the name. 287 288 If filename is of the form "foo.zip+bar.zip", unzip foo.zip into a 289 temp dir, then unzip bar.zip into that_dir/BOOTABLE_IMAGES. 290 291 Returns (tempdir, zipobj) where zipobj is a zipfile.ZipFile (of the 292 main file), open for reading. 293 """ 294 295 tmp = tempfile.mkdtemp(prefix="targetfiles-") 296 OPTIONS.tempfiles.append(tmp) 297 298 def unzip_to_dir(filename, dirname): 299 cmd = ["unzip", "-o", "-q", filename, "-d", dirname] 300 if pattern is not None: 301 cmd.append(pattern) 302 p = Run(cmd, stdout=subprocess.PIPE) 303 p.communicate() 304 if p.returncode != 0: 305 raise ExternalError("failed to unzip input target-files \"%s\"" % 306 (filename,)) 307 308 m = re.match(r"^(.*[.]zip)\+(.*[.]zip)$", filename, re.IGNORECASE) 309 if m: 310 unzip_to_dir(m.group(1), tmp) 311 unzip_to_dir(m.group(2), os.path.join(tmp, "BOOTABLE_IMAGES")) 312 filename = m.group(1) 313 else: 314 unzip_to_dir(filename, tmp) 315 316 return tmp, zipfile.ZipFile(filename, "r") 317 318 319def GetKeyPasswords(keylist): 320 """Given a list of keys, prompt the user to enter passwords for 321 those which require them. Return a {key: password} dict. password 322 will be None if the key has no password.""" 323 324 no_passwords = [] 325 need_passwords = [] 326 devnull = open("/dev/null", "w+b") 327 for k in sorted(keylist): 328 # We don't need a password for things that aren't really keys. 329 if k in SPECIAL_CERT_STRINGS: 330 no_passwords.append(k) 331 continue 332 333 p = Run(["openssl", "pkcs8", "-in", k+".pk8", 334 "-inform", "DER", "-nocrypt"], 335 stdin=devnull.fileno(), 336 stdout=devnull.fileno(), 337 stderr=subprocess.STDOUT) 338 p.communicate() 339 if p.returncode == 0: 340 no_passwords.append(k) 341 else: 342 need_passwords.append(k) 343 devnull.close() 344 345 key_passwords = PasswordManager().GetPasswords(need_passwords) 346 key_passwords.update(dict.fromkeys(no_passwords, None)) 347 return key_passwords 348 349 350def SignFile(input_name, output_name, key, password, align=None, 351 whole_file=False): 352 """Sign the input_name zip/jar/apk, producing output_name. Use the 353 given key and password (the latter may be None if the key does not 354 have a password. 355 356 If align is an integer > 1, zipalign is run to align stored files in 357 the output zip on 'align'-byte boundaries. 358 359 If whole_file is true, use the "-w" option to SignApk to embed a 360 signature that covers the whole file in the archive comment of the 361 zip file. 362 """ 363 364 if align == 0 or align == 1: 365 align = None 366 367 if align: 368 temp = tempfile.NamedTemporaryFile() 369 sign_name = temp.name 370 else: 371 sign_name = output_name 372 373 cmd = ["java", "-Xmx2048m", "-jar", 374 os.path.join(OPTIONS.search_path, "framework", "signapk.jar")] 375 if whole_file: 376 cmd.append("-w") 377 cmd.extend([key + ".x509.pem", key + ".pk8", 378 input_name, sign_name]) 379 380 p = Run(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE) 381 if password is not None: 382 password += "\n" 383 p.communicate(password) 384 if p.returncode != 0: 385 raise ExternalError("signapk.jar failed: return code %s" % (p.returncode,)) 386 387 if align: 388 p = Run(["zipalign", "-f", str(align), sign_name, output_name]) 389 p.communicate() 390 if p.returncode != 0: 391 raise ExternalError("zipalign failed: return code %s" % (p.returncode,)) 392 temp.close() 393 394 395def CheckSize(data, target, info_dict): 396 """Check the data string passed against the max size limit, if 397 any, for the given target. Raise exception if the data is too big. 398 Print a warning if the data is nearing the maximum size.""" 399 400 if target.endswith(".img"): target = target[:-4] 401 mount_point = "/" + target 402 403 if info_dict["fstab"]: 404 if mount_point == "/userdata": mount_point = "/data" 405 p = info_dict["fstab"][mount_point] 406 fs_type = p.fs_type 407 device = p.device 408 if "/" in device: 409 device = device[device.rfind("/")+1:] 410 limit = info_dict.get(device + "_size", None) 411 if not fs_type or not limit: return 412 413 if fs_type == "yaffs2": 414 # image size should be increased by 1/64th to account for the 415 # spare area (64 bytes per 2k page) 416 limit = limit / 2048 * (2048+64) 417 size = len(data) 418 pct = float(size) * 100.0 / limit 419 msg = "%s size (%d) is %.2f%% of limit (%d)" % (target, size, pct, limit) 420 if pct >= 99.0: 421 raise ExternalError(msg) 422 elif pct >= 95.0: 423 print 424 print " WARNING: ", msg 425 print 426 elif OPTIONS.verbose: 427 print " ", msg 428 429 430def ReadApkCerts(tf_zip): 431 """Given a target_files ZipFile, parse the META/apkcerts.txt file 432 and return a {package: cert} dict.""" 433 certmap = {} 434 for line in tf_zip.read("META/apkcerts.txt").split("\n"): 435 line = line.strip() 436 if not line: continue 437 m = re.match(r'^name="(.*)"\s+certificate="(.*)"\s+' 438 r'private_key="(.*)"$', line) 439 if m: 440 name, cert, privkey = m.groups() 441 if cert in SPECIAL_CERT_STRINGS and not privkey: 442 certmap[name] = cert 443 elif (cert.endswith(".x509.pem") and 444 privkey.endswith(".pk8") and 445 cert[:-9] == privkey[:-4]): 446 certmap[name] = cert[:-9] 447 else: 448 raise ValueError("failed to parse line from apkcerts.txt:\n" + line) 449 return certmap 450 451 452COMMON_DOCSTRING = """ 453 -p (--path) <dir> 454 Prepend <dir>/bin to the list of places to search for binaries 455 run by this script, and expect to find jars in <dir>/framework. 456 457 -s (--device_specific) <file> 458 Path to the python module containing device-specific 459 releasetools code. 460 461 -x (--extra) <key=value> 462 Add a key/value pair to the 'extras' dict, which device-specific 463 extension code may look at. 464 465 -v (--verbose) 466 Show command lines being executed. 467 468 -h (--help) 469 Display this usage message and exit. 470""" 471 472def Usage(docstring): 473 print docstring.rstrip("\n") 474 print COMMON_DOCSTRING 475 476 477def ParseOptions(argv, 478 docstring, 479 extra_opts="", extra_long_opts=(), 480 extra_option_handler=None): 481 """Parse the options in argv and return any arguments that aren't 482 flags. docstring is the calling module's docstring, to be displayed 483 for errors and -h. extra_opts and extra_long_opts are for flags 484 defined by the caller, which are processed by passing them to 485 extra_option_handler.""" 486 487 try: 488 opts, args = getopt.getopt( 489 argv, "hvp:s:x:" + extra_opts, 490 ["help", "verbose", "path=", "device_specific=", "extra="] + 491 list(extra_long_opts)) 492 except getopt.GetoptError, err: 493 Usage(docstring) 494 print "**", str(err), "**" 495 sys.exit(2) 496 497 path_specified = False 498 499 for o, a in opts: 500 if o in ("-h", "--help"): 501 Usage(docstring) 502 sys.exit() 503 elif o in ("-v", "--verbose"): 504 OPTIONS.verbose = True 505 elif o in ("-p", "--path"): 506 OPTIONS.search_path = a 507 elif o in ("-s", "--device_specific"): 508 OPTIONS.device_specific = a 509 elif o in ("-x", "--extra"): 510 key, value = a.split("=", 1) 511 OPTIONS.extras[key] = value 512 else: 513 if extra_option_handler is None or not extra_option_handler(o, a): 514 assert False, "unknown option \"%s\"" % (o,) 515 516 os.environ["PATH"] = (os.path.join(OPTIONS.search_path, "bin") + 517 os.pathsep + os.environ["PATH"]) 518 519 return args 520 521 522def Cleanup(): 523 for i in OPTIONS.tempfiles: 524 if os.path.isdir(i): 525 shutil.rmtree(i) 526 else: 527 os.remove(i) 528 529 530class PasswordManager(object): 531 def __init__(self): 532 self.editor = os.getenv("EDITOR", None) 533 self.pwfile = os.getenv("ANDROID_PW_FILE", None) 534 535 def GetPasswords(self, items): 536 """Get passwords corresponding to each string in 'items', 537 returning a dict. (The dict may have keys in addition to the 538 values in 'items'.) 539 540 Uses the passwords in $ANDROID_PW_FILE if available, letting the 541 user edit that file to add more needed passwords. If no editor is 542 available, or $ANDROID_PW_FILE isn't define, prompts the user 543 interactively in the ordinary way. 544 """ 545 546 current = self.ReadFile() 547 548 first = True 549 while True: 550 missing = [] 551 for i in items: 552 if i not in current or not current[i]: 553 missing.append(i) 554 # Are all the passwords already in the file? 555 if not missing: return current 556 557 for i in missing: 558 current[i] = "" 559 560 if not first: 561 print "key file %s still missing some passwords." % (self.pwfile,) 562 answer = raw_input("try to edit again? [y]> ").strip() 563 if answer and answer[0] not in 'yY': 564 raise RuntimeError("key passwords unavailable") 565 first = False 566 567 current = self.UpdateAndReadFile(current) 568 569 def PromptResult(self, current): 570 """Prompt the user to enter a value (password) for each key in 571 'current' whose value is fales. Returns a new dict with all the 572 values. 573 """ 574 result = {} 575 for k, v in sorted(current.iteritems()): 576 if v: 577 result[k] = v 578 else: 579 while True: 580 result[k] = getpass.getpass("Enter password for %s key> " 581 % (k,)).strip() 582 if result[k]: break 583 return result 584 585 def UpdateAndReadFile(self, current): 586 if not self.editor or not self.pwfile: 587 return self.PromptResult(current) 588 589 f = open(self.pwfile, "w") 590 os.chmod(self.pwfile, 0600) 591 f.write("# Enter key passwords between the [[[ ]]] brackets.\n") 592 f.write("# (Additional spaces are harmless.)\n\n") 593 594 first_line = None 595 sorted = [(not v, k, v) for (k, v) in current.iteritems()] 596 sorted.sort() 597 for i, (_, k, v) in enumerate(sorted): 598 f.write("[[[ %s ]]] %s\n" % (v, k)) 599 if not v and first_line is None: 600 # position cursor on first line with no password. 601 first_line = i + 4 602 f.close() 603 604 p = Run([self.editor, "+%d" % (first_line,), self.pwfile]) 605 _, _ = p.communicate() 606 607 return self.ReadFile() 608 609 def ReadFile(self): 610 result = {} 611 if self.pwfile is None: return result 612 try: 613 f = open(self.pwfile, "r") 614 for line in f: 615 line = line.strip() 616 if not line or line[0] == '#': continue 617 m = re.match(r"^\[\[\[\s*(.*?)\s*\]\]\]\s*(\S+)$", line) 618 if not m: 619 print "failed to parse password file: ", line 620 else: 621 result[m.group(2)] = m.group(1) 622 f.close() 623 except IOError, e: 624 if e.errno != errno.ENOENT: 625 print "error reading password file: ", str(e) 626 return result 627 628 629def ZipWriteStr(zip, filename, data, perms=0644): 630 # use a fixed timestamp so the output is repeatable. 631 zinfo = zipfile.ZipInfo(filename=filename, 632 date_time=(2009, 1, 1, 0, 0, 0)) 633 zinfo.compress_type = zip.compression 634 zinfo.external_attr = perms << 16 635 zip.writestr(zinfo, data) 636 637 638class DeviceSpecificParams(object): 639 module = None 640 def __init__(self, **kwargs): 641 """Keyword arguments to the constructor become attributes of this 642 object, which is passed to all functions in the device-specific 643 module.""" 644 for k, v in kwargs.iteritems(): 645 setattr(self, k, v) 646 self.extras = OPTIONS.extras 647 648 if self.module is None: 649 path = OPTIONS.device_specific 650 if not path: return 651 try: 652 if os.path.isdir(path): 653 info = imp.find_module("releasetools", [path]) 654 else: 655 d, f = os.path.split(path) 656 b, x = os.path.splitext(f) 657 if x == ".py": 658 f = b 659 info = imp.find_module(f, [d]) 660 self.module = imp.load_module("device_specific", *info) 661 except ImportError: 662 print "unable to load device-specific module; assuming none" 663 664 def _DoCall(self, function_name, *args, **kwargs): 665 """Call the named function in the device-specific module, passing 666 the given args and kwargs. The first argument to the call will be 667 the DeviceSpecific object itself. If there is no module, or the 668 module does not define the function, return the value of the 669 'default' kwarg (which itself defaults to None).""" 670 if self.module is None or not hasattr(self.module, function_name): 671 return kwargs.get("default", None) 672 return getattr(self.module, function_name)(*((self,) + args), **kwargs) 673 674 def FullOTA_Assertions(self): 675 """Called after emitting the block of assertions at the top of a 676 full OTA package. Implementations can add whatever additional 677 assertions they like.""" 678 return self._DoCall("FullOTA_Assertions") 679 680 def FullOTA_InstallBegin(self): 681 """Called at the start of full OTA installation.""" 682 return self._DoCall("FullOTA_InstallBegin") 683 684 def FullOTA_InstallEnd(self): 685 """Called at the end of full OTA installation; typically this is 686 used to install the image for the device's baseband processor.""" 687 return self._DoCall("FullOTA_InstallEnd") 688 689 def IncrementalOTA_Assertions(self): 690 """Called after emitting the block of assertions at the top of an 691 incremental OTA package. Implementations can add whatever 692 additional assertions they like.""" 693 return self._DoCall("IncrementalOTA_Assertions") 694 695 def IncrementalOTA_VerifyBegin(self): 696 """Called at the start of the verification phase of incremental 697 OTA installation; additional checks can be placed here to abort 698 the script before any changes are made.""" 699 return self._DoCall("IncrementalOTA_VerifyBegin") 700 701 def IncrementalOTA_VerifyEnd(self): 702 """Called at the end of the verification phase of incremental OTA 703 installation; additional checks can be placed here to abort the 704 script before any changes are made.""" 705 return self._DoCall("IncrementalOTA_VerifyEnd") 706 707 def IncrementalOTA_InstallBegin(self): 708 """Called at the start of incremental OTA installation (after 709 verification is complete).""" 710 return self._DoCall("IncrementalOTA_InstallBegin") 711 712 def IncrementalOTA_InstallEnd(self): 713 """Called at the end of incremental OTA installation; typically 714 this is used to install the image for the device's baseband 715 processor.""" 716 return self._DoCall("IncrementalOTA_InstallEnd") 717 718class File(object): 719 def __init__(self, name, data): 720 self.name = name 721 self.data = data 722 self.size = len(data) 723 self.sha1 = sha1(data).hexdigest() 724 725 @classmethod 726 def FromLocalFile(cls, name, diskname): 727 f = open(diskname, "rb") 728 data = f.read() 729 f.close() 730 return File(name, data) 731 732 def WriteToTemp(self): 733 t = tempfile.NamedTemporaryFile() 734 t.write(self.data) 735 t.flush() 736 return t 737 738 def AddToZip(self, z): 739 ZipWriteStr(z, self.name, self.data) 740 741DIFF_PROGRAM_BY_EXT = { 742 ".gz" : "imgdiff", 743 ".zip" : ["imgdiff", "-z"], 744 ".jar" : ["imgdiff", "-z"], 745 ".apk" : ["imgdiff", "-z"], 746 ".img" : "imgdiff", 747 } 748 749class Difference(object): 750 def __init__(self, tf, sf, diff_program=None): 751 self.tf = tf 752 self.sf = sf 753 self.patch = None 754 self.diff_program = diff_program 755 756 def ComputePatch(self): 757 """Compute the patch (as a string of data) needed to turn sf into 758 tf. Returns the same tuple as GetPatch().""" 759 760 tf = self.tf 761 sf = self.sf 762 763 if self.diff_program: 764 diff_program = self.diff_program 765 else: 766 ext = os.path.splitext(tf.name)[1] 767 diff_program = DIFF_PROGRAM_BY_EXT.get(ext, "bsdiff") 768 769 ttemp = tf.WriteToTemp() 770 stemp = sf.WriteToTemp() 771 772 ext = os.path.splitext(tf.name)[1] 773 774 try: 775 ptemp = tempfile.NamedTemporaryFile() 776 if isinstance(diff_program, list): 777 cmd = copy.copy(diff_program) 778 else: 779 cmd = [diff_program] 780 cmd.append(stemp.name) 781 cmd.append(ttemp.name) 782 cmd.append(ptemp.name) 783 p = Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) 784 _, err = p.communicate() 785 if err or p.returncode != 0: 786 print "WARNING: failure running %s:\n%s\n" % (diff_program, err) 787 return None 788 diff = ptemp.read() 789 finally: 790 ptemp.close() 791 stemp.close() 792 ttemp.close() 793 794 self.patch = diff 795 return self.tf, self.sf, self.patch 796 797 798 def GetPatch(self): 799 """Return a tuple (target_file, source_file, patch_data). 800 patch_data may be None if ComputePatch hasn't been called, or if 801 computing the patch failed.""" 802 return self.tf, self.sf, self.patch 803 804 805def ComputeDifferences(diffs): 806 """Call ComputePatch on all the Difference objects in 'diffs'.""" 807 print len(diffs), "diffs to compute" 808 809 # Do the largest files first, to try and reduce the long-pole effect. 810 by_size = [(i.tf.size, i) for i in diffs] 811 by_size.sort(reverse=True) 812 by_size = [i[1] for i in by_size] 813 814 lock = threading.Lock() 815 diff_iter = iter(by_size) # accessed under lock 816 817 def worker(): 818 try: 819 lock.acquire() 820 for d in diff_iter: 821 lock.release() 822 start = time.time() 823 d.ComputePatch() 824 dur = time.time() - start 825 lock.acquire() 826 827 tf, sf, patch = d.GetPatch() 828 if sf.name == tf.name: 829 name = tf.name 830 else: 831 name = "%s (%s)" % (tf.name, sf.name) 832 if patch is None: 833 print "patching failed! %s" % (name,) 834 else: 835 print "%8.2f sec %8d / %8d bytes (%6.2f%%) %s" % ( 836 dur, len(patch), tf.size, 100.0 * len(patch) / tf.size, name) 837 lock.release() 838 except Exception, e: 839 print e 840 raise 841 842 # start worker threads; wait for them all to finish. 843 threads = [threading.Thread(target=worker) 844 for i in range(OPTIONS.worker_threads)] 845 for th in threads: 846 th.start() 847 while threads: 848 threads.pop().join() 849 850 851# map recovery.fstab's fs_types to mount/format "partition types" 852PARTITION_TYPES = { "yaffs2": "MTD", "mtd": "MTD", 853 "ext4": "EMMC", "emmc": "EMMC" } 854 855def GetTypeAndDevice(mount_point, info): 856 fstab = info["fstab"] 857 if fstab: 858 return PARTITION_TYPES[fstab[mount_point].fs_type], fstab[mount_point].device 859 else: 860 return None 861