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