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