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