common.py revision f5085a8f94c50689d6ab522f2c5b9c4ae01d323d
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    data = BuildBootableImage(os.path.join(unpack_dir, tree_subdir),
346                              os.path.join(unpack_dir, fs_config),
347                              info_dict)
348    if data:
349      return File(name, data)
350    return None
351
352
353def UnzipTemp(filename, pattern=None):
354  """Unzip the given archive into a temporary directory and return the name.
355
356  If filename is of the form "foo.zip+bar.zip", unzip foo.zip into a
357  temp dir, then unzip bar.zip into that_dir/BOOTABLE_IMAGES.
358
359  Returns (tempdir, zipobj) where zipobj is a zipfile.ZipFile (of the
360  main file), open for reading.
361  """
362
363  tmp = tempfile.mkdtemp(prefix="targetfiles-")
364  OPTIONS.tempfiles.append(tmp)
365
366  def unzip_to_dir(filename, dirname):
367    cmd = ["unzip", "-o", "-q", filename, "-d", dirname]
368    if pattern is not None:
369      cmd.append(pattern)
370    p = Run(cmd, stdout=subprocess.PIPE)
371    p.communicate()
372    if p.returncode != 0:
373      raise ExternalError("failed to unzip input target-files \"%s\"" %
374                          (filename,))
375
376  m = re.match(r"^(.*[.]zip)\+(.*[.]zip)$", filename, re.IGNORECASE)
377  if m:
378    unzip_to_dir(m.group(1), tmp)
379    unzip_to_dir(m.group(2), os.path.join(tmp, "BOOTABLE_IMAGES"))
380    filename = m.group(1)
381  else:
382    unzip_to_dir(filename, tmp)
383
384  return tmp, zipfile.ZipFile(filename, "r")
385
386
387def GetKeyPasswords(keylist):
388  """Given a list of keys, prompt the user to enter passwords for
389  those which require them.  Return a {key: password} dict.  password
390  will be None if the key has no password."""
391
392  no_passwords = []
393  need_passwords = []
394  key_passwords = {}
395  devnull = open("/dev/null", "w+b")
396  for k in sorted(keylist):
397    # We don't need a password for things that aren't really keys.
398    if k in SPECIAL_CERT_STRINGS:
399      no_passwords.append(k)
400      continue
401
402    p = Run(["openssl", "pkcs8", "-in", k+OPTIONS.private_key_suffix,
403             "-inform", "DER", "-nocrypt"],
404            stdin=devnull.fileno(),
405            stdout=devnull.fileno(),
406            stderr=subprocess.STDOUT)
407    p.communicate()
408    if p.returncode == 0:
409      # Definitely an unencrypted key.
410      no_passwords.append(k)
411    else:
412      p = Run(["openssl", "pkcs8", "-in", k+OPTIONS.private_key_suffix,
413               "-inform", "DER", "-passin", "pass:"],
414              stdin=devnull.fileno(),
415              stdout=devnull.fileno(),
416              stderr=subprocess.PIPE)
417      stdout, stderr = p.communicate()
418      if p.returncode == 0:
419        # Encrypted key with empty string as password.
420        key_passwords[k] = ''
421      elif stderr.startswith('Error decrypting key'):
422        # Definitely encrypted key.
423        # It would have said "Error reading key" if it didn't parse correctly.
424        need_passwords.append(k)
425      else:
426        # Potentially, a type of key that openssl doesn't understand.
427        # We'll let the routines in signapk.jar handle it.
428        no_passwords.append(k)
429  devnull.close()
430
431  key_passwords.update(PasswordManager().GetPasswords(need_passwords))
432  key_passwords.update(dict.fromkeys(no_passwords, None))
433  return key_passwords
434
435
436def SignFile(input_name, output_name, key, password, align=None,
437             whole_file=False):
438  """Sign the input_name zip/jar/apk, producing output_name.  Use the
439  given key and password (the latter may be None if the key does not
440  have a password.
441
442  If align is an integer > 1, zipalign is run to align stored files in
443  the output zip on 'align'-byte boundaries.
444
445  If whole_file is true, use the "-w" option to SignApk to embed a
446  signature that covers the whole file in the archive comment of the
447  zip file.
448  """
449
450  if align == 0 or align == 1:
451    align = None
452
453  if align:
454    temp = tempfile.NamedTemporaryFile()
455    sign_name = temp.name
456  else:
457    sign_name = output_name
458
459  cmd = [OPTIONS.java_path, "-Xmx2048m", "-jar",
460         os.path.join(OPTIONS.search_path, OPTIONS.signapk_path)]
461  cmd.extend(OPTIONS.extra_signapk_args)
462  if whole_file:
463    cmd.append("-w")
464  cmd.extend([key + OPTIONS.public_key_suffix,
465              key + OPTIONS.private_key_suffix,
466              input_name, sign_name])
467
468  p = Run(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
469  if password is not None:
470    password += "\n"
471  p.communicate(password)
472  if p.returncode != 0:
473    raise ExternalError("signapk.jar failed: return code %s" % (p.returncode,))
474
475  if align:
476    p = Run(["zipalign", "-f", str(align), sign_name, output_name])
477    p.communicate()
478    if p.returncode != 0:
479      raise ExternalError("zipalign failed: return code %s" % (p.returncode,))
480    temp.close()
481
482
483def CheckSize(data, target, info_dict):
484  """Check the data string passed against the max size limit, if
485  any, for the given target.  Raise exception if the data is too big.
486  Print a warning if the data is nearing the maximum size."""
487
488  if target.endswith(".img"): target = target[:-4]
489  mount_point = "/" + target
490
491  fs_type = None
492  limit = None
493  if info_dict["fstab"]:
494    if mount_point == "/userdata": mount_point = "/data"
495    p = info_dict["fstab"][mount_point]
496    fs_type = p.fs_type
497    device = p.device
498    if "/" in device:
499      device = device[device.rfind("/")+1:]
500    limit = info_dict.get(device + "_size", None)
501  if not fs_type or not limit: return
502
503  if fs_type == "yaffs2":
504    # image size should be increased by 1/64th to account for the
505    # spare area (64 bytes per 2k page)
506    limit = limit / 2048 * (2048+64)
507  size = len(data)
508  pct = float(size) * 100.0 / limit
509  msg = "%s size (%d) is %.2f%% of limit (%d)" % (target, size, pct, limit)
510  if pct >= 99.0:
511    raise ExternalError(msg)
512  elif pct >= 95.0:
513    print
514    print "  WARNING: ", msg
515    print
516  elif OPTIONS.verbose:
517    print "  ", msg
518
519
520def ReadApkCerts(tf_zip):
521  """Given a target_files ZipFile, parse the META/apkcerts.txt file
522  and return a {package: cert} dict."""
523  certmap = {}
524  for line in tf_zip.read("META/apkcerts.txt").split("\n"):
525    line = line.strip()
526    if not line: continue
527    m = re.match(r'^name="(.*)"\s+certificate="(.*)"\s+'
528                 r'private_key="(.*)"$', line)
529    if m:
530      name, cert, privkey = m.groups()
531      public_key_suffix_len = len(OPTIONS.public_key_suffix)
532      private_key_suffix_len = len(OPTIONS.private_key_suffix)
533      if cert in SPECIAL_CERT_STRINGS and not privkey:
534        certmap[name] = cert
535      elif (cert.endswith(OPTIONS.public_key_suffix) and
536            privkey.endswith(OPTIONS.private_key_suffix) and
537            cert[:-public_key_suffix_len] == privkey[:-private_key_suffix_len]):
538        certmap[name] = cert[:-public_key_suffix_len]
539      else:
540        raise ValueError("failed to parse line from apkcerts.txt:\n" + line)
541  return certmap
542
543
544COMMON_DOCSTRING = """
545  -p  (--path)  <dir>
546      Prepend <dir>/bin to the list of places to search for binaries
547      run by this script, and expect to find jars in <dir>/framework.
548
549  -s  (--device_specific) <file>
550      Path to the python module containing device-specific
551      releasetools code.
552
553  -x  (--extra)  <key=value>
554      Add a key/value pair to the 'extras' dict, which device-specific
555      extension code may look at.
556
557  -v  (--verbose)
558      Show command lines being executed.
559
560  -h  (--help)
561      Display this usage message and exit.
562"""
563
564def Usage(docstring):
565  print docstring.rstrip("\n")
566  print COMMON_DOCSTRING
567
568
569def ParseOptions(argv,
570                 docstring,
571                 extra_opts="", extra_long_opts=(),
572                 extra_option_handler=None):
573  """Parse the options in argv and return any arguments that aren't
574  flags.  docstring is the calling module's docstring, to be displayed
575  for errors and -h.  extra_opts and extra_long_opts are for flags
576  defined by the caller, which are processed by passing them to
577  extra_option_handler."""
578
579  try:
580    opts, args = getopt.getopt(
581        argv, "hvp:s:x:" + extra_opts,
582        ["help", "verbose", "path=", "signapk_path=", "extra_signapk_args=",
583         "java_path=", "java_args=", "public_key_suffix=",
584         "private_key_suffix=", "device_specific=", "extra="] +
585        list(extra_long_opts))
586  except getopt.GetoptError, err:
587    Usage(docstring)
588    print "**", str(err), "**"
589    sys.exit(2)
590
591  path_specified = False
592
593  for o, a in opts:
594    if o in ("-h", "--help"):
595      Usage(docstring)
596      sys.exit()
597    elif o in ("-v", "--verbose"):
598      OPTIONS.verbose = True
599    elif o in ("-p", "--path"):
600      OPTIONS.search_path = a
601    elif o in ("--signapk_path",):
602      OPTIONS.signapk_path = a
603    elif o in ("--extra_signapk_args",):
604      OPTIONS.extra_signapk_args = shlex.split(a)
605    elif o in ("--java_path",):
606      OPTIONS.java_path = a
607    elif o in ("--public_key_suffix",):
608      OPTIONS.public_key_suffix = a
609    elif o in ("--private_key_suffix",):
610      OPTIONS.private_key_suffix = a
611    elif o in ("-s", "--device_specific"):
612      OPTIONS.device_specific = a
613    elif o in ("-x", "--extra"):
614      key, value = a.split("=", 1)
615      OPTIONS.extras[key] = value
616    else:
617      if extra_option_handler is None or not extra_option_handler(o, a):
618        assert False, "unknown option \"%s\"" % (o,)
619
620  os.environ["PATH"] = (os.path.join(OPTIONS.search_path, "bin") +
621                        os.pathsep + os.environ["PATH"])
622
623  return args
624
625
626def Cleanup():
627  for i in OPTIONS.tempfiles:
628    if os.path.isdir(i):
629      shutil.rmtree(i)
630    else:
631      os.remove(i)
632
633
634class PasswordManager(object):
635  def __init__(self):
636    self.editor = os.getenv("EDITOR", None)
637    self.pwfile = os.getenv("ANDROID_PW_FILE", None)
638
639  def GetPasswords(self, items):
640    """Get passwords corresponding to each string in 'items',
641    returning a dict.  (The dict may have keys in addition to the
642    values in 'items'.)
643
644    Uses the passwords in $ANDROID_PW_FILE if available, letting the
645    user edit that file to add more needed passwords.  If no editor is
646    available, or $ANDROID_PW_FILE isn't define, prompts the user
647    interactively in the ordinary way.
648    """
649
650    current = self.ReadFile()
651
652    first = True
653    while True:
654      missing = []
655      for i in items:
656        if i not in current or not current[i]:
657          missing.append(i)
658      # Are all the passwords already in the file?
659      if not missing: return current
660
661      for i in missing:
662        current[i] = ""
663
664      if not first:
665        print "key file %s still missing some passwords." % (self.pwfile,)
666        answer = raw_input("try to edit again? [y]> ").strip()
667        if answer and answer[0] not in 'yY':
668          raise RuntimeError("key passwords unavailable")
669      first = False
670
671      current = self.UpdateAndReadFile(current)
672
673  def PromptResult(self, current):
674    """Prompt the user to enter a value (password) for each key in
675    'current' whose value is fales.  Returns a new dict with all the
676    values.
677    """
678    result = {}
679    for k, v in sorted(current.iteritems()):
680      if v:
681        result[k] = v
682      else:
683        while True:
684          result[k] = getpass.getpass("Enter password for %s key> "
685                                      % (k,)).strip()
686          if result[k]: break
687    return result
688
689  def UpdateAndReadFile(self, current):
690    if not self.editor or not self.pwfile:
691      return self.PromptResult(current)
692
693    f = open(self.pwfile, "w")
694    os.chmod(self.pwfile, 0600)
695    f.write("# Enter key passwords between the [[[ ]]] brackets.\n")
696    f.write("# (Additional spaces are harmless.)\n\n")
697
698    first_line = None
699    sorted = [(not v, k, v) for (k, v) in current.iteritems()]
700    sorted.sort()
701    for i, (_, k, v) in enumerate(sorted):
702      f.write("[[[  %s  ]]] %s\n" % (v, k))
703      if not v and first_line is None:
704        # position cursor on first line with no password.
705        first_line = i + 4
706    f.close()
707
708    p = Run([self.editor, "+%d" % (first_line,), self.pwfile])
709    _, _ = p.communicate()
710
711    return self.ReadFile()
712
713  def ReadFile(self):
714    result = {}
715    if self.pwfile is None: return result
716    try:
717      f = open(self.pwfile, "r")
718      for line in f:
719        line = line.strip()
720        if not line or line[0] == '#': continue
721        m = re.match(r"^\[\[\[\s*(.*?)\s*\]\]\]\s*(\S+)$", line)
722        if not m:
723          print "failed to parse password file: ", line
724        else:
725          result[m.group(2)] = m.group(1)
726      f.close()
727    except IOError, e:
728      if e.errno != errno.ENOENT:
729        print "error reading password file: ", str(e)
730    return result
731
732
733def ZipWriteStr(zip, filename, data, perms=0644):
734  # use a fixed timestamp so the output is repeatable.
735  zinfo = zipfile.ZipInfo(filename=filename,
736                          date_time=(2009, 1, 1, 0, 0, 0))
737  zinfo.compress_type = zip.compression
738  zinfo.external_attr = perms << 16
739  zip.writestr(zinfo, data)
740
741
742class DeviceSpecificParams(object):
743  module = None
744  def __init__(self, **kwargs):
745    """Keyword arguments to the constructor become attributes of this
746    object, which is passed to all functions in the device-specific
747    module."""
748    for k, v in kwargs.iteritems():
749      setattr(self, k, v)
750    self.extras = OPTIONS.extras
751
752    if self.module is None:
753      path = OPTIONS.device_specific
754      if not path: return
755      try:
756        if os.path.isdir(path):
757          info = imp.find_module("releasetools", [path])
758        else:
759          d, f = os.path.split(path)
760          b, x = os.path.splitext(f)
761          if x == ".py":
762            f = b
763          info = imp.find_module(f, [d])
764        self.module = imp.load_module("device_specific", *info)
765      except ImportError:
766        print "unable to load device-specific module; assuming none"
767
768  def _DoCall(self, function_name, *args, **kwargs):
769    """Call the named function in the device-specific module, passing
770    the given args and kwargs.  The first argument to the call will be
771    the DeviceSpecific object itself.  If there is no module, or the
772    module does not define the function, return the value of the
773    'default' kwarg (which itself defaults to None)."""
774    if self.module is None or not hasattr(self.module, function_name):
775      return kwargs.get("default", None)
776    return getattr(self.module, function_name)(*((self,) + args), **kwargs)
777
778  def FullOTA_Assertions(self):
779    """Called after emitting the block of assertions at the top of a
780    full OTA package.  Implementations can add whatever additional
781    assertions they like."""
782    return self._DoCall("FullOTA_Assertions")
783
784  def FullOTA_InstallBegin(self):
785    """Called at the start of full OTA installation."""
786    return self._DoCall("FullOTA_InstallBegin")
787
788  def FullOTA_InstallEnd(self):
789    """Called at the end of full OTA installation; typically this is
790    used to install the image for the device's baseband processor."""
791    return self._DoCall("FullOTA_InstallEnd")
792
793  def IncrementalOTA_Assertions(self):
794    """Called after emitting the block of assertions at the top of an
795    incremental OTA package.  Implementations can add whatever
796    additional assertions they like."""
797    return self._DoCall("IncrementalOTA_Assertions")
798
799  def IncrementalOTA_VerifyBegin(self):
800    """Called at the start of the verification phase of incremental
801    OTA installation; additional checks can be placed here to abort
802    the script before any changes are made."""
803    return self._DoCall("IncrementalOTA_VerifyBegin")
804
805  def IncrementalOTA_VerifyEnd(self):
806    """Called at the end of the verification phase of incremental OTA
807    installation; additional checks can be placed here to abort the
808    script before any changes are made."""
809    return self._DoCall("IncrementalOTA_VerifyEnd")
810
811  def IncrementalOTA_InstallBegin(self):
812    """Called at the start of incremental OTA installation (after
813    verification is complete)."""
814    return self._DoCall("IncrementalOTA_InstallBegin")
815
816  def IncrementalOTA_InstallEnd(self):
817    """Called at the end of incremental OTA installation; typically
818    this is used to install the image for the device's baseband
819    processor."""
820    return self._DoCall("IncrementalOTA_InstallEnd")
821
822class File(object):
823  def __init__(self, name, data):
824    self.name = name
825    self.data = data
826    self.size = len(data)
827    self.sha1 = sha1(data).hexdigest()
828
829  @classmethod
830  def FromLocalFile(cls, name, diskname):
831    f = open(diskname, "rb")
832    data = f.read()
833    f.close()
834    return File(name, data)
835
836  def WriteToTemp(self):
837    t = tempfile.NamedTemporaryFile()
838    t.write(self.data)
839    t.flush()
840    return t
841
842  def AddToZip(self, z):
843    ZipWriteStr(z, self.name, self.data)
844
845DIFF_PROGRAM_BY_EXT = {
846    ".gz" : "imgdiff",
847    ".zip" : ["imgdiff", "-z"],
848    ".jar" : ["imgdiff", "-z"],
849    ".apk" : ["imgdiff", "-z"],
850    ".img" : "imgdiff",
851    }
852
853class Difference(object):
854  def __init__(self, tf, sf, diff_program=None):
855    self.tf = tf
856    self.sf = sf
857    self.patch = None
858    self.diff_program = diff_program
859
860  def ComputePatch(self):
861    """Compute the patch (as a string of data) needed to turn sf into
862    tf.  Returns the same tuple as GetPatch()."""
863
864    tf = self.tf
865    sf = self.sf
866
867    if self.diff_program:
868      diff_program = self.diff_program
869    else:
870      ext = os.path.splitext(tf.name)[1]
871      diff_program = DIFF_PROGRAM_BY_EXT.get(ext, "bsdiff")
872
873    ttemp = tf.WriteToTemp()
874    stemp = sf.WriteToTemp()
875
876    ext = os.path.splitext(tf.name)[1]
877
878    try:
879      ptemp = tempfile.NamedTemporaryFile()
880      if isinstance(diff_program, list):
881        cmd = copy.copy(diff_program)
882      else:
883        cmd = [diff_program]
884      cmd.append(stemp.name)
885      cmd.append(ttemp.name)
886      cmd.append(ptemp.name)
887      p = Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
888      _, err = p.communicate()
889      if err or p.returncode != 0:
890        print "WARNING: failure running %s:\n%s\n" % (diff_program, err)
891        return None
892      diff = ptemp.read()
893    finally:
894      ptemp.close()
895      stemp.close()
896      ttemp.close()
897
898    self.patch = diff
899    return self.tf, self.sf, self.patch
900
901
902  def GetPatch(self):
903    """Return a tuple (target_file, source_file, patch_data).
904    patch_data may be None if ComputePatch hasn't been called, or if
905    computing the patch failed."""
906    return self.tf, self.sf, self.patch
907
908
909def ComputeDifferences(diffs):
910  """Call ComputePatch on all the Difference objects in 'diffs'."""
911  print len(diffs), "diffs to compute"
912
913  # Do the largest files first, to try and reduce the long-pole effect.
914  by_size = [(i.tf.size, i) for i in diffs]
915  by_size.sort(reverse=True)
916  by_size = [i[1] for i in by_size]
917
918  lock = threading.Lock()
919  diff_iter = iter(by_size)   # accessed under lock
920
921  def worker():
922    try:
923      lock.acquire()
924      for d in diff_iter:
925        lock.release()
926        start = time.time()
927        d.ComputePatch()
928        dur = time.time() - start
929        lock.acquire()
930
931        tf, sf, patch = d.GetPatch()
932        if sf.name == tf.name:
933          name = tf.name
934        else:
935          name = "%s (%s)" % (tf.name, sf.name)
936        if patch is None:
937          print "patching failed!                                  %s" % (name,)
938        else:
939          print "%8.2f sec %8d / %8d bytes (%6.2f%%) %s" % (
940              dur, len(patch), tf.size, 100.0 * len(patch) / tf.size, name)
941      lock.release()
942    except Exception, e:
943      print e
944      raise
945
946  # start worker threads; wait for them all to finish.
947  threads = [threading.Thread(target=worker)
948             for i in range(OPTIONS.worker_threads)]
949  for th in threads:
950    th.start()
951  while threads:
952    threads.pop().join()
953
954
955# map recovery.fstab's fs_types to mount/format "partition types"
956PARTITION_TYPES = { "yaffs2": "MTD", "mtd": "MTD",
957                    "ext4": "EMMC", "emmc": "EMMC" }
958
959def GetTypeAndDevice(mount_point, info):
960  fstab = info["fstab"]
961  if fstab:
962    return PARTITION_TYPES[fstab[mount_point].fs_type], fstab[mount_point].device
963  else:
964    return None
965
966
967def ParseCertificate(data):
968  """Parse a PEM-format certificate."""
969  cert = []
970  save = False
971  for line in data.split("\n"):
972    if "--END CERTIFICATE--" in line:
973      break
974    if save:
975      cert.append(line)
976    if "--BEGIN CERTIFICATE--" in line:
977      save = True
978  cert = "".join(cert).decode('base64')
979  return cert
980