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