common.py revision fffe1d5e9c65dcb2321643e25d77d465f7763ba6
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):
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  ramdisk_img = tempfile.NamedTemporaryFile()
210  img = tempfile.NamedTemporaryFile()
211
212  if os.access(fs_config_file, os.F_OK):
213    cmd = ["mkbootfs", "-f", fs_config_file, os.path.join(sourcedir, "RAMDISK")]
214  else:
215    cmd = ["mkbootfs", os.path.join(sourcedir, "RAMDISK")]
216  p1 = Run(cmd, stdout=subprocess.PIPE)
217  p2 = Run(["minigzip"],
218           stdin=p1.stdout, stdout=ramdisk_img.file.fileno())
219
220  p2.wait()
221  p1.wait()
222  assert p1.returncode == 0, "mkbootfs of %s ramdisk failed" % (targetname,)
223  assert p2.returncode == 0, "minigzip of %s ramdisk failed" % (targetname,)
224
225  cmd = ["mkbootimg", "--kernel", os.path.join(sourcedir, "kernel")]
226
227  fn = os.path.join(sourcedir, "cmdline")
228  if os.access(fn, os.F_OK):
229    cmd.append("--cmdline")
230    cmd.append(open(fn).read().rstrip("\n"))
231
232  fn = os.path.join(sourcedir, "base")
233  if os.access(fn, os.F_OK):
234    cmd.append("--base")
235    cmd.append(open(fn).read().rstrip("\n"))
236
237  fn = os.path.join(sourcedir, "pagesize")
238  if os.access(fn, os.F_OK):
239    cmd.append("--pagesize")
240    cmd.append(open(fn).read().rstrip("\n"))
241
242  cmd.extend(["--ramdisk", ramdisk_img.name,
243              "--output", img.name])
244
245  p = Run(cmd, stdout=subprocess.PIPE)
246  p.communicate()
247  assert p.returncode == 0, "mkbootimg of %s image failed" % (
248      os.path.basename(sourcedir),)
249
250  img.seek(os.SEEK_SET, 0)
251  data = img.read()
252
253  ramdisk_img.close()
254  img.close()
255
256  return data
257
258
259def GetBootableImage(name, prebuilt_name, unpack_dir, tree_subdir):
260  """Return a File object (with name 'name') with the desired bootable
261  image.  Look for it in 'unpack_dir'/BOOTABLE_IMAGES under the name
262  'prebuilt_name', otherwise construct it from the source files in
263  'unpack_dir'/'tree_subdir'."""
264
265  prebuilt_path = os.path.join(unpack_dir, "BOOTABLE_IMAGES", prebuilt_name)
266  if os.path.exists(prebuilt_path):
267    print "using prebuilt %s..." % (prebuilt_name,)
268    return File.FromLocalFile(name, prebuilt_path)
269  else:
270    print "building image from target_files %s..." % (tree_subdir,)
271    fs_config = "META/" + tree_subdir.lower() + "_filesystem_config.txt"
272    return File(name, BuildBootableImage(os.path.join(unpack_dir, tree_subdir),
273                                         os.path.join(unpack_dir, fs_config)))
274
275
276def UnzipTemp(filename, pattern=None):
277  """Unzip the given archive into a temporary directory and return the name.
278
279  If filename is of the form "foo.zip+bar.zip", unzip foo.zip into a
280  temp dir, then unzip bar.zip into that_dir/BOOTABLE_IMAGES.
281
282  Returns (tempdir, zipobj) where zipobj is a zipfile.ZipFile (of the
283  main file), open for reading.
284  """
285
286  tmp = tempfile.mkdtemp(prefix="targetfiles-")
287  OPTIONS.tempfiles.append(tmp)
288
289  def unzip_to_dir(filename, dirname):
290    cmd = ["unzip", "-o", "-q", filename, "-d", dirname]
291    if pattern is not None:
292      cmd.append(pattern)
293    p = Run(cmd, stdout=subprocess.PIPE)
294    p.communicate()
295    if p.returncode != 0:
296      raise ExternalError("failed to unzip input target-files \"%s\"" %
297                          (filename,))
298
299  m = re.match(r"^(.*[.]zip)\+(.*[.]zip)$", filename, re.IGNORECASE)
300  if m:
301    unzip_to_dir(m.group(1), tmp)
302    unzip_to_dir(m.group(2), os.path.join(tmp, "BOOTABLE_IMAGES"))
303    filename = m.group(1)
304  else:
305    unzip_to_dir(filename, tmp)
306
307  return tmp, zipfile.ZipFile(filename, "r")
308
309
310def GetKeyPasswords(keylist):
311  """Given a list of keys, prompt the user to enter passwords for
312  those which require them.  Return a {key: password} dict.  password
313  will be None if the key has no password."""
314
315  no_passwords = []
316  need_passwords = []
317  devnull = open("/dev/null", "w+b")
318  for k in sorted(keylist):
319    # We don't need a password for things that aren't really keys.
320    if k in SPECIAL_CERT_STRINGS:
321      no_passwords.append(k)
322      continue
323
324    p = Run(["openssl", "pkcs8", "-in", k+".pk8",
325             "-inform", "DER", "-nocrypt"],
326            stdin=devnull.fileno(),
327            stdout=devnull.fileno(),
328            stderr=subprocess.STDOUT)
329    p.communicate()
330    if p.returncode == 0:
331      no_passwords.append(k)
332    else:
333      need_passwords.append(k)
334  devnull.close()
335
336  key_passwords = PasswordManager().GetPasswords(need_passwords)
337  key_passwords.update(dict.fromkeys(no_passwords, None))
338  return key_passwords
339
340
341def SignFile(input_name, output_name, key, password, align=None,
342             whole_file=False):
343  """Sign the input_name zip/jar/apk, producing output_name.  Use the
344  given key and password (the latter may be None if the key does not
345  have a password.
346
347  If align is an integer > 1, zipalign is run to align stored files in
348  the output zip on 'align'-byte boundaries.
349
350  If whole_file is true, use the "-w" option to SignApk to embed a
351  signature that covers the whole file in the archive comment of the
352  zip file.
353  """
354
355  if align == 0 or align == 1:
356    align = None
357
358  if align:
359    temp = tempfile.NamedTemporaryFile()
360    sign_name = temp.name
361  else:
362    sign_name = output_name
363
364  cmd = ["java", "-Xmx2048m", "-jar",
365           os.path.join(OPTIONS.search_path, "framework", "signapk.jar")]
366  if whole_file:
367    cmd.append("-w")
368  cmd.extend([key + ".x509.pem", key + ".pk8",
369              input_name, sign_name])
370
371  p = Run(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
372  if password is not None:
373    password += "\n"
374  p.communicate(password)
375  if p.returncode != 0:
376    raise ExternalError("signapk.jar failed: return code %s" % (p.returncode,))
377
378  if align:
379    p = Run(["zipalign", "-f", str(align), sign_name, output_name])
380    p.communicate()
381    if p.returncode != 0:
382      raise ExternalError("zipalign failed: return code %s" % (p.returncode,))
383    temp.close()
384
385
386def CheckSize(data, target, info_dict):
387  """Check the data string passed against the max size limit, if
388  any, for the given target.  Raise exception if the data is too big.
389  Print a warning if the data is nearing the maximum size."""
390
391  if target.endswith(".img"): target = target[:-4]
392  mount_point = "/" + target
393
394  if info_dict["fstab"]:
395    if mount_point == "/userdata": mount_point = "/data"
396    p = info_dict["fstab"][mount_point]
397    fs_type = p.fs_type
398    device = p.device
399    if "/" in device:
400      device = device[device.rfind("/")+1:]
401    limit = info_dict.get(device + "_size", None)
402  if not fs_type or not limit: return
403
404  if fs_type == "yaffs2":
405    # image size should be increased by 1/64th to account for the
406    # spare area (64 bytes per 2k page)
407    limit = limit / 2048 * (2048+64)
408  size = len(data)
409  pct = float(size) * 100.0 / limit
410  msg = "%s size (%d) is %.2f%% of limit (%d)" % (target, size, pct, limit)
411  if pct >= 99.0:
412    raise ExternalError(msg)
413  elif pct >= 95.0:
414    print
415    print "  WARNING: ", msg
416    print
417  elif OPTIONS.verbose:
418    print "  ", msg
419
420
421def ReadApkCerts(tf_zip):
422  """Given a target_files ZipFile, parse the META/apkcerts.txt file
423  and return a {package: cert} dict."""
424  certmap = {}
425  for line in tf_zip.read("META/apkcerts.txt").split("\n"):
426    line = line.strip()
427    if not line: continue
428    m = re.match(r'^name="(.*)"\s+certificate="(.*)"\s+'
429                 r'private_key="(.*)"$', line)
430    if m:
431      name, cert, privkey = m.groups()
432      if cert in SPECIAL_CERT_STRINGS and not privkey:
433        certmap[name] = cert
434      elif (cert.endswith(".x509.pem") and
435            privkey.endswith(".pk8") and
436            cert[:-9] == privkey[:-4]):
437        certmap[name] = cert[:-9]
438      else:
439        raise ValueError("failed to parse line from apkcerts.txt:\n" + line)
440  return certmap
441
442
443COMMON_DOCSTRING = """
444  -p  (--path)  <dir>
445      Prepend <dir>/bin to the list of places to search for binaries
446      run by this script, and expect to find jars in <dir>/framework.
447
448  -s  (--device_specific) <file>
449      Path to the python module containing device-specific
450      releasetools code.
451
452  -x  (--extra)  <key=value>
453      Add a key/value pair to the 'extras' dict, which device-specific
454      extension code may look at.
455
456  -v  (--verbose)
457      Show command lines being executed.
458
459  -h  (--help)
460      Display this usage message and exit.
461"""
462
463def Usage(docstring):
464  print docstring.rstrip("\n")
465  print COMMON_DOCSTRING
466
467
468def ParseOptions(argv,
469                 docstring,
470                 extra_opts="", extra_long_opts=(),
471                 extra_option_handler=None):
472  """Parse the options in argv and return any arguments that aren't
473  flags.  docstring is the calling module's docstring, to be displayed
474  for errors and -h.  extra_opts and extra_long_opts are for flags
475  defined by the caller, which are processed by passing them to
476  extra_option_handler."""
477
478  try:
479    opts, args = getopt.getopt(
480        argv, "hvp:s:x:" + extra_opts,
481        ["help", "verbose", "path=", "device_specific=", "extra="] +
482          list(extra_long_opts))
483  except getopt.GetoptError, err:
484    Usage(docstring)
485    print "**", str(err), "**"
486    sys.exit(2)
487
488  path_specified = False
489
490  for o, a in opts:
491    if o in ("-h", "--help"):
492      Usage(docstring)
493      sys.exit()
494    elif o in ("-v", "--verbose"):
495      OPTIONS.verbose = True
496    elif o in ("-p", "--path"):
497      OPTIONS.search_path = a
498    elif o in ("-s", "--device_specific"):
499      OPTIONS.device_specific = a
500    elif o in ("-x", "--extra"):
501      key, value = a.split("=", 1)
502      OPTIONS.extras[key] = value
503    else:
504      if extra_option_handler is None or not extra_option_handler(o, a):
505        assert False, "unknown option \"%s\"" % (o,)
506
507  os.environ["PATH"] = (os.path.join(OPTIONS.search_path, "bin") +
508                        os.pathsep + os.environ["PATH"])
509
510  return args
511
512
513def Cleanup():
514  for i in OPTIONS.tempfiles:
515    if os.path.isdir(i):
516      shutil.rmtree(i)
517    else:
518      os.remove(i)
519
520
521class PasswordManager(object):
522  def __init__(self):
523    self.editor = os.getenv("EDITOR", None)
524    self.pwfile = os.getenv("ANDROID_PW_FILE", None)
525
526  def GetPasswords(self, items):
527    """Get passwords corresponding to each string in 'items',
528    returning a dict.  (The dict may have keys in addition to the
529    values in 'items'.)
530
531    Uses the passwords in $ANDROID_PW_FILE if available, letting the
532    user edit that file to add more needed passwords.  If no editor is
533    available, or $ANDROID_PW_FILE isn't define, prompts the user
534    interactively in the ordinary way.
535    """
536
537    current = self.ReadFile()
538
539    first = True
540    while True:
541      missing = []
542      for i in items:
543        if i not in current or not current[i]:
544          missing.append(i)
545      # Are all the passwords already in the file?
546      if not missing: return current
547
548      for i in missing:
549        current[i] = ""
550
551      if not first:
552        print "key file %s still missing some passwords." % (self.pwfile,)
553        answer = raw_input("try to edit again? [y]> ").strip()
554        if answer and answer[0] not in 'yY':
555          raise RuntimeError("key passwords unavailable")
556      first = False
557
558      current = self.UpdateAndReadFile(current)
559
560  def PromptResult(self, current):
561    """Prompt the user to enter a value (password) for each key in
562    'current' whose value is fales.  Returns a new dict with all the
563    values.
564    """
565    result = {}
566    for k, v in sorted(current.iteritems()):
567      if v:
568        result[k] = v
569      else:
570        while True:
571          result[k] = getpass.getpass("Enter password for %s key> "
572                                      % (k,)).strip()
573          if result[k]: break
574    return result
575
576  def UpdateAndReadFile(self, current):
577    if not self.editor or not self.pwfile:
578      return self.PromptResult(current)
579
580    f = open(self.pwfile, "w")
581    os.chmod(self.pwfile, 0600)
582    f.write("# Enter key passwords between the [[[ ]]] brackets.\n")
583    f.write("# (Additional spaces are harmless.)\n\n")
584
585    first_line = None
586    sorted = [(not v, k, v) for (k, v) in current.iteritems()]
587    sorted.sort()
588    for i, (_, k, v) in enumerate(sorted):
589      f.write("[[[  %s  ]]] %s\n" % (v, k))
590      if not v and first_line is None:
591        # position cursor on first line with no password.
592        first_line = i + 4
593    f.close()
594
595    p = Run([self.editor, "+%d" % (first_line,), self.pwfile])
596    _, _ = p.communicate()
597
598    return self.ReadFile()
599
600  def ReadFile(self):
601    result = {}
602    if self.pwfile is None: return result
603    try:
604      f = open(self.pwfile, "r")
605      for line in f:
606        line = line.strip()
607        if not line or line[0] == '#': continue
608        m = re.match(r"^\[\[\[\s*(.*?)\s*\]\]\]\s*(\S+)$", line)
609        if not m:
610          print "failed to parse password file: ", line
611        else:
612          result[m.group(2)] = m.group(1)
613      f.close()
614    except IOError, e:
615      if e.errno != errno.ENOENT:
616        print "error reading password file: ", str(e)
617    return result
618
619
620def ZipWriteStr(zip, filename, data, perms=0644):
621  # use a fixed timestamp so the output is repeatable.
622  zinfo = zipfile.ZipInfo(filename=filename,
623                          date_time=(2009, 1, 1, 0, 0, 0))
624  zinfo.compress_type = zip.compression
625  zinfo.external_attr = perms << 16
626  zip.writestr(zinfo, data)
627
628
629class DeviceSpecificParams(object):
630  module = None
631  def __init__(self, **kwargs):
632    """Keyword arguments to the constructor become attributes of this
633    object, which is passed to all functions in the device-specific
634    module."""
635    for k, v in kwargs.iteritems():
636      setattr(self, k, v)
637    self.extras = OPTIONS.extras
638
639    if self.module is None:
640      path = OPTIONS.device_specific
641      if not path: return
642      try:
643        if os.path.isdir(path):
644          info = imp.find_module("releasetools", [path])
645        else:
646          d, f = os.path.split(path)
647          b, x = os.path.splitext(f)
648          if x == ".py":
649            f = b
650          info = imp.find_module(f, [d])
651        self.module = imp.load_module("device_specific", *info)
652      except ImportError:
653        print "unable to load device-specific module; assuming none"
654
655  def _DoCall(self, function_name, *args, **kwargs):
656    """Call the named function in the device-specific module, passing
657    the given args and kwargs.  The first argument to the call will be
658    the DeviceSpecific object itself.  If there is no module, or the
659    module does not define the function, return the value of the
660    'default' kwarg (which itself defaults to None)."""
661    if self.module is None or not hasattr(self.module, function_name):
662      return kwargs.get("default", None)
663    return getattr(self.module, function_name)(*((self,) + args), **kwargs)
664
665  def FullOTA_Assertions(self):
666    """Called after emitting the block of assertions at the top of a
667    full OTA package.  Implementations can add whatever additional
668    assertions they like."""
669    return self._DoCall("FullOTA_Assertions")
670
671  def FullOTA_InstallBegin(self):
672    """Called at the start of full OTA installation."""
673    return self._DoCall("FullOTA_InstallBegin")
674
675  def FullOTA_InstallEnd(self):
676    """Called at the end of full OTA installation; typically this is
677    used to install the image for the device's baseband processor."""
678    return self._DoCall("FullOTA_InstallEnd")
679
680  def IncrementalOTA_Assertions(self):
681    """Called after emitting the block of assertions at the top of an
682    incremental OTA package.  Implementations can add whatever
683    additional assertions they like."""
684    return self._DoCall("IncrementalOTA_Assertions")
685
686  def IncrementalOTA_VerifyBegin(self):
687    """Called at the start of the verification phase of incremental
688    OTA installation; additional checks can be placed here to abort
689    the script before any changes are made."""
690    return self._DoCall("IncrementalOTA_VerifyBegin")
691
692  def IncrementalOTA_VerifyEnd(self):
693    """Called at the end of the verification phase of incremental OTA
694    installation; additional checks can be placed here to abort the
695    script before any changes are made."""
696    return self._DoCall("IncrementalOTA_VerifyEnd")
697
698  def IncrementalOTA_InstallBegin(self):
699    """Called at the start of incremental OTA installation (after
700    verification is complete)."""
701    return self._DoCall("IncrementalOTA_InstallBegin")
702
703  def IncrementalOTA_InstallEnd(self):
704    """Called at the end of incremental OTA installation; typically
705    this is used to install the image for the device's baseband
706    processor."""
707    return self._DoCall("IncrementalOTA_InstallEnd")
708
709class File(object):
710  def __init__(self, name, data):
711    self.name = name
712    self.data = data
713    self.size = len(data)
714    self.sha1 = sha1(data).hexdigest()
715
716  @classmethod
717  def FromLocalFile(cls, name, diskname):
718    f = open(diskname, "rb")
719    data = f.read()
720    f.close()
721    return File(name, data)
722
723  def WriteToTemp(self):
724    t = tempfile.NamedTemporaryFile()
725    t.write(self.data)
726    t.flush()
727    return t
728
729  def AddToZip(self, z):
730    ZipWriteStr(z, self.name, self.data)
731
732DIFF_PROGRAM_BY_EXT = {
733    ".gz" : "imgdiff",
734    ".zip" : ["imgdiff", "-z"],
735    ".jar" : ["imgdiff", "-z"],
736    ".apk" : ["imgdiff", "-z"],
737    ".img" : "imgdiff",
738    }
739
740class Difference(object):
741  def __init__(self, tf, sf):
742    self.tf = tf
743    self.sf = sf
744    self.patch = None
745
746  def ComputePatch(self):
747    """Compute the patch (as a string of data) needed to turn sf into
748    tf.  Returns the same tuple as GetPatch()."""
749
750    tf = self.tf
751    sf = self.sf
752
753    ext = os.path.splitext(tf.name)[1]
754    diff_program = DIFF_PROGRAM_BY_EXT.get(ext, "bsdiff")
755
756    ttemp = tf.WriteToTemp()
757    stemp = sf.WriteToTemp()
758
759    ext = os.path.splitext(tf.name)[1]
760
761    try:
762      ptemp = tempfile.NamedTemporaryFile()
763      if isinstance(diff_program, list):
764        cmd = copy.copy(diff_program)
765      else:
766        cmd = [diff_program]
767      cmd.append(stemp.name)
768      cmd.append(ttemp.name)
769      cmd.append(ptemp.name)
770      p = Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
771      _, err = p.communicate()
772      if err or p.returncode != 0:
773        print "WARNING: failure running %s:\n%s\n" % (diff_program, err)
774        return None
775      diff = ptemp.read()
776    finally:
777      ptemp.close()
778      stemp.close()
779      ttemp.close()
780
781    self.patch = diff
782    return self.tf, self.sf, self.patch
783
784
785  def GetPatch(self):
786    """Return a tuple (target_file, source_file, patch_data).
787    patch_data may be None if ComputePatch hasn't been called, or if
788    computing the patch failed."""
789    return self.tf, self.sf, self.patch
790
791
792def ComputeDifferences(diffs):
793  """Call ComputePatch on all the Difference objects in 'diffs'."""
794  print len(diffs), "diffs to compute"
795
796  # Do the largest files first, to try and reduce the long-pole effect.
797  by_size = [(i.tf.size, i) for i in diffs]
798  by_size.sort(reverse=True)
799  by_size = [i[1] for i in by_size]
800
801  lock = threading.Lock()
802  diff_iter = iter(by_size)   # accessed under lock
803
804  def worker():
805    try:
806      lock.acquire()
807      for d in diff_iter:
808        lock.release()
809        start = time.time()
810        d.ComputePatch()
811        dur = time.time() - start
812        lock.acquire()
813
814        tf, sf, patch = d.GetPatch()
815        if sf.name == tf.name:
816          name = tf.name
817        else:
818          name = "%s (%s)" % (tf.name, sf.name)
819        if patch is None:
820          print "patching failed!                                  %s" % (name,)
821        else:
822          print "%8.2f sec %8d / %8d bytes (%6.2f%%) %s" % (
823              dur, len(patch), tf.size, 100.0 * len(patch) / tf.size, name)
824      lock.release()
825    except Exception, e:
826      print e
827      raise
828
829  # start worker threads; wait for them all to finish.
830  threads = [threading.Thread(target=worker)
831             for i in range(OPTIONS.worker_threads)]
832  for th in threads:
833    th.start()
834  while threads:
835    threads.pop().join()
836
837
838# map recovery.fstab's fs_types to mount/format "partition types"
839PARTITION_TYPES = { "yaffs2": "MTD", "mtd": "MTD",
840                    "ext4": "EMMC", "emmc": "EMMC" }
841
842def GetTypeAndDevice(mount_point, info):
843  fstab = info["fstab"]
844  if fstab:
845    return PARTITION_TYPES[fstab[mount_point].fs_type], fstab[mount_point].device
846  else:
847    return None
848