common.py revision 94919cd385327e250b4c9340c715fce981dddbcc
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 Cleanup():
656  for i in OPTIONS.tempfiles:
657    if os.path.isdir(i):
658      shutil.rmtree(i)
659    else:
660      os.remove(i)
661
662
663class PasswordManager(object):
664  def __init__(self):
665    self.editor = os.getenv("EDITOR", None)
666    self.pwfile = os.getenv("ANDROID_PW_FILE", None)
667
668  def GetPasswords(self, items):
669    """Get passwords corresponding to each string in 'items',
670    returning a dict.  (The dict may have keys in addition to the
671    values in 'items'.)
672
673    Uses the passwords in $ANDROID_PW_FILE if available, letting the
674    user edit that file to add more needed passwords.  If no editor is
675    available, or $ANDROID_PW_FILE isn't define, prompts the user
676    interactively in the ordinary way.
677    """
678
679    current = self.ReadFile()
680
681    first = True
682    while True:
683      missing = []
684      for i in items:
685        if i not in current or not current[i]:
686          missing.append(i)
687      # Are all the passwords already in the file?
688      if not missing: return current
689
690      for i in missing:
691        current[i] = ""
692
693      if not first:
694        print "key file %s still missing some passwords." % (self.pwfile,)
695        answer = raw_input("try to edit again? [y]> ").strip()
696        if answer and answer[0] not in 'yY':
697          raise RuntimeError("key passwords unavailable")
698      first = False
699
700      current = self.UpdateAndReadFile(current)
701
702  def PromptResult(self, current):
703    """Prompt the user to enter a value (password) for each key in
704    'current' whose value is fales.  Returns a new dict with all the
705    values.
706    """
707    result = {}
708    for k, v in sorted(current.iteritems()):
709      if v:
710        result[k] = v
711      else:
712        while True:
713          result[k] = getpass.getpass("Enter password for %s key> "
714                                      % (k,)).strip()
715          if result[k]: break
716    return result
717
718  def UpdateAndReadFile(self, current):
719    if not self.editor or not self.pwfile:
720      return self.PromptResult(current)
721
722    f = open(self.pwfile, "w")
723    os.chmod(self.pwfile, 0600)
724    f.write("# Enter key passwords between the [[[ ]]] brackets.\n")
725    f.write("# (Additional spaces are harmless.)\n\n")
726
727    first_line = None
728    sorted = [(not v, k, v) for (k, v) in current.iteritems()]
729    sorted.sort()
730    for i, (_, k, v) in enumerate(sorted):
731      f.write("[[[  %s  ]]] %s\n" % (v, k))
732      if not v and first_line is None:
733        # position cursor on first line with no password.
734        first_line = i + 4
735    f.close()
736
737    p = Run([self.editor, "+%d" % (first_line,), self.pwfile])
738    _, _ = p.communicate()
739
740    return self.ReadFile()
741
742  def ReadFile(self):
743    result = {}
744    if self.pwfile is None: return result
745    try:
746      f = open(self.pwfile, "r")
747      for line in f:
748        line = line.strip()
749        if not line or line[0] == '#': continue
750        m = re.match(r"^\[\[\[\s*(.*?)\s*\]\]\]\s*(\S+)$", line)
751        if not m:
752          print "failed to parse password file: ", line
753        else:
754          result[m.group(2)] = m.group(1)
755      f.close()
756    except IOError, e:
757      if e.errno != errno.ENOENT:
758        print "error reading password file: ", str(e)
759    return result
760
761
762def ZipWriteStr(zip, filename, data, perms=0644, compression=None):
763  # use a fixed timestamp so the output is repeatable.
764  zinfo = zipfile.ZipInfo(filename=filename,
765                          date_time=(2009, 1, 1, 0, 0, 0))
766  if compression is None:
767    zinfo.compress_type = zip.compression
768  else:
769    zinfo.compress_type = compression
770  zinfo.external_attr = perms << 16
771  zip.writestr(zinfo, data)
772
773
774class DeviceSpecificParams(object):
775  module = None
776  def __init__(self, **kwargs):
777    """Keyword arguments to the constructor become attributes of this
778    object, which is passed to all functions in the device-specific
779    module."""
780    for k, v in kwargs.iteritems():
781      setattr(self, k, v)
782    self.extras = OPTIONS.extras
783
784    if self.module is None:
785      path = OPTIONS.device_specific
786      if not path: return
787      try:
788        if os.path.isdir(path):
789          info = imp.find_module("releasetools", [path])
790        else:
791          d, f = os.path.split(path)
792          b, x = os.path.splitext(f)
793          if x == ".py":
794            f = b
795          info = imp.find_module(f, [d])
796        print "loaded device-specific extensions from", path
797        self.module = imp.load_module("device_specific", *info)
798      except ImportError:
799        print "unable to load device-specific module; assuming none"
800
801  def _DoCall(self, function_name, *args, **kwargs):
802    """Call the named function in the device-specific module, passing
803    the given args and kwargs.  The first argument to the call will be
804    the DeviceSpecific object itself.  If there is no module, or the
805    module does not define the function, return the value of the
806    'default' kwarg (which itself defaults to None)."""
807    if self.module is None or not hasattr(self.module, function_name):
808      return kwargs.get("default", None)
809    return getattr(self.module, function_name)(*((self,) + args), **kwargs)
810
811  def FullOTA_Assertions(self):
812    """Called after emitting the block of assertions at the top of a
813    full OTA package.  Implementations can add whatever additional
814    assertions they like."""
815    return self._DoCall("FullOTA_Assertions")
816
817  def FullOTA_InstallBegin(self):
818    """Called at the start of full OTA installation."""
819    return self._DoCall("FullOTA_InstallBegin")
820
821  def FullOTA_InstallEnd(self):
822    """Called at the end of full OTA installation; typically this is
823    used to install the image for the device's baseband processor."""
824    return self._DoCall("FullOTA_InstallEnd")
825
826  def IncrementalOTA_Assertions(self):
827    """Called after emitting the block of assertions at the top of an
828    incremental OTA package.  Implementations can add whatever
829    additional assertions they like."""
830    return self._DoCall("IncrementalOTA_Assertions")
831
832  def IncrementalOTA_VerifyBegin(self):
833    """Called at the start of the verification phase of incremental
834    OTA installation; additional checks can be placed here to abort
835    the script before any changes are made."""
836    return self._DoCall("IncrementalOTA_VerifyBegin")
837
838  def IncrementalOTA_VerifyEnd(self):
839    """Called at the end of the verification phase of incremental OTA
840    installation; additional checks can be placed here to abort the
841    script before any changes are made."""
842    return self._DoCall("IncrementalOTA_VerifyEnd")
843
844  def IncrementalOTA_InstallBegin(self):
845    """Called at the start of incremental OTA installation (after
846    verification is complete)."""
847    return self._DoCall("IncrementalOTA_InstallBegin")
848
849  def IncrementalOTA_InstallEnd(self):
850    """Called at the end of incremental OTA installation; typically
851    this is used to install the image for the device's baseband
852    processor."""
853    return self._DoCall("IncrementalOTA_InstallEnd")
854
855class File(object):
856  def __init__(self, name, data):
857    self.name = name
858    self.data = data
859    self.size = len(data)
860    self.sha1 = sha1(data).hexdigest()
861
862  @classmethod
863  def FromLocalFile(cls, name, diskname):
864    f = open(diskname, "rb")
865    data = f.read()
866    f.close()
867    return File(name, data)
868
869  def WriteToTemp(self):
870    t = tempfile.NamedTemporaryFile()
871    t.write(self.data)
872    t.flush()
873    return t
874
875  def AddToZip(self, z, compression=None):
876    ZipWriteStr(z, self.name, self.data, compression=compression)
877
878DIFF_PROGRAM_BY_EXT = {
879    ".gz" : "imgdiff",
880    ".zip" : ["imgdiff", "-z"],
881    ".jar" : ["imgdiff", "-z"],
882    ".apk" : ["imgdiff", "-z"],
883    ".img" : "imgdiff",
884    }
885
886class Difference(object):
887  def __init__(self, tf, sf, diff_program=None):
888    self.tf = tf
889    self.sf = sf
890    self.patch = None
891    self.diff_program = diff_program
892
893  def ComputePatch(self):
894    """Compute the patch (as a string of data) needed to turn sf into
895    tf.  Returns the same tuple as GetPatch()."""
896
897    tf = self.tf
898    sf = self.sf
899
900    if self.diff_program:
901      diff_program = self.diff_program
902    else:
903      ext = os.path.splitext(tf.name)[1]
904      diff_program = DIFF_PROGRAM_BY_EXT.get(ext, "bsdiff")
905
906    ttemp = tf.WriteToTemp()
907    stemp = sf.WriteToTemp()
908
909    ext = os.path.splitext(tf.name)[1]
910
911    try:
912      ptemp = tempfile.NamedTemporaryFile()
913      if isinstance(diff_program, list):
914        cmd = copy.copy(diff_program)
915      else:
916        cmd = [diff_program]
917      cmd.append(stemp.name)
918      cmd.append(ttemp.name)
919      cmd.append(ptemp.name)
920      p = Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
921      err = []
922      def run():
923        _, e = p.communicate()
924        if e: err.append(e)
925      th = threading.Thread(target=run)
926      th.start()
927      th.join(timeout=300)   # 5 mins
928      if th.is_alive():
929        print "WARNING: diff command timed out"
930        p.terminate()
931        th.join(5)
932        if th.is_alive():
933          p.kill()
934          th.join()
935
936      if err or p.returncode != 0:
937        print "WARNING: failure running %s:\n%s\n" % (
938            diff_program, "".join(err))
939        self.patch = None
940        return None, None, None
941      diff = ptemp.read()
942    finally:
943      ptemp.close()
944      stemp.close()
945      ttemp.close()
946
947    self.patch = diff
948    return self.tf, self.sf, self.patch
949
950
951  def GetPatch(self):
952    """Return a tuple (target_file, source_file, patch_data).
953    patch_data may be None if ComputePatch hasn't been called, or if
954    computing the patch failed."""
955    return self.tf, self.sf, self.patch
956
957
958def ComputeDifferences(diffs):
959  """Call ComputePatch on all the Difference objects in 'diffs'."""
960  print len(diffs), "diffs to compute"
961
962  # Do the largest files first, to try and reduce the long-pole effect.
963  by_size = [(i.tf.size, i) for i in diffs]
964  by_size.sort(reverse=True)
965  by_size = [i[1] for i in by_size]
966
967  lock = threading.Lock()
968  diff_iter = iter(by_size)   # accessed under lock
969
970  def worker():
971    try:
972      lock.acquire()
973      for d in diff_iter:
974        lock.release()
975        start = time.time()
976        d.ComputePatch()
977        dur = time.time() - start
978        lock.acquire()
979
980        tf, sf, patch = d.GetPatch()
981        if sf.name == tf.name:
982          name = tf.name
983        else:
984          name = "%s (%s)" % (tf.name, sf.name)
985        if patch is None:
986          print "patching failed!                                  %s" % (name,)
987        else:
988          print "%8.2f sec %8d / %8d bytes (%6.2f%%) %s" % (
989              dur, len(patch), tf.size, 100.0 * len(patch) / tf.size, name)
990      lock.release()
991    except Exception, e:
992      print e
993      raise
994
995  # start worker threads; wait for them all to finish.
996  threads = [threading.Thread(target=worker)
997             for i in range(OPTIONS.worker_threads)]
998  for th in threads:
999    th.start()
1000  while threads:
1001    threads.pop().join()
1002
1003
1004# map recovery.fstab's fs_types to mount/format "partition types"
1005PARTITION_TYPES = { "yaffs2": "MTD", "mtd": "MTD",
1006                    "ext4": "EMMC", "emmc": "EMMC",
1007                    "f2fs": "EMMC" }
1008
1009def GetTypeAndDevice(mount_point, info):
1010  fstab = info["fstab"]
1011  if fstab:
1012    return PARTITION_TYPES[fstab[mount_point].fs_type], fstab[mount_point].device
1013  else:
1014    return None
1015
1016
1017def ParseCertificate(data):
1018  """Parse a PEM-format certificate."""
1019  cert = []
1020  save = False
1021  for line in data.split("\n"):
1022    if "--END CERTIFICATE--" in line:
1023      break
1024    if save:
1025      cert.append(line)
1026    if "--BEGIN CERTIFICATE--" in line:
1027      save = True
1028  cert = "".join(cert).decode('base64')
1029  return cert
1030
1031def XDelta3(source_path, target_path, output_path):
1032  diff_program = ["xdelta3", "-0", "-B", str(64<<20), "-e", "-f", "-s"]
1033  diff_program.append(source_path)
1034  diff_program.append(target_path)
1035  diff_program.append(output_path)
1036  p = Run(diff_program, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
1037  p.communicate()
1038  assert p.returncode == 0, "Couldn't produce patch"
1039
1040def XZ(path):
1041  compress_program = ["xz", "-zk", "-9", "--check=crc32"]
1042  compress_program.append(path)
1043  p = Run(compress_program, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
1044  p.communicate()
1045  assert p.returncode == 0, "Couldn't compress patch"
1046
1047def MakePartitionPatch(source_file, target_file, partition):
1048  with tempfile.NamedTemporaryFile() as output_file:
1049    XDelta3(source_file.name, target_file.name, output_file.name)
1050    XZ(output_file.name)
1051    with open(output_file.name + ".xz") as patch_file:
1052      patch_data = patch_file.read()
1053      os.unlink(patch_file.name)
1054      return File(partition + ".muimg.p", patch_data)
1055
1056def MakeRecoveryPatch(input_dir, output_sink, recovery_img, boot_img,
1057                      info_dict=None):
1058  """Generate a binary patch that creates the recovery image starting
1059  with the boot image.  (Most of the space in these images is just the
1060  kernel, which is identical for the two, so the resulting patch
1061  should be efficient.)  Add it to the output zip, along with a shell
1062  script that is run from init.rc on first boot to actually do the
1063  patching and install the new recovery image.
1064
1065  recovery_img and boot_img should be File objects for the
1066  corresponding images.  info should be the dictionary returned by
1067  common.LoadInfoDict() on the input target_files.
1068  """
1069
1070  if info_dict is None:
1071    info_dict = OPTIONS.info_dict
1072
1073  diff_program = ["imgdiff"]
1074  path = os.path.join(input_dir, "SYSTEM", "etc", "recovery-resource.dat")
1075  if os.path.exists(path):
1076    diff_program.append("-b")
1077    diff_program.append(path)
1078    bonus_args = "-b /system/etc/recovery-resource.dat"
1079  else:
1080    bonus_args = ""
1081
1082  d = Difference(recovery_img, boot_img, diff_program=diff_program)
1083  _, _, patch = d.ComputePatch()
1084  output_sink("recovery-from-boot.p", patch)
1085
1086  td_pair = GetTypeAndDevice("/boot", info_dict)
1087  if not td_pair:
1088    return
1089  boot_type, boot_device = td_pair
1090  td_pair = GetTypeAndDevice("/recovery", info_dict)
1091  if not td_pair:
1092    return
1093  recovery_type, recovery_device = td_pair
1094
1095  sh = """#!/system/bin/sh
1096if ! applypatch -c %(recovery_type)s:%(recovery_device)s:%(recovery_size)d:%(recovery_sha1)s; then
1097  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"
1098else
1099  log -t recovery "Recovery image already installed"
1100fi
1101""" % { 'boot_size': boot_img.size,
1102        'boot_sha1': boot_img.sha1,
1103        'recovery_size': recovery_img.size,
1104        'recovery_sha1': recovery_img.sha1,
1105        'boot_type': boot_type,
1106        'boot_device': boot_device,
1107        'recovery_type': recovery_type,
1108        'recovery_device': recovery_device,
1109        'bonus_args': bonus_args,
1110        }
1111
1112  # The install script location moved from /system/etc to /system/bin
1113  # in the L release.  Parse the init.rc file to find out where the
1114  # target-files expects it to be, and put it there.
1115  sh_location = "etc/install-recovery.sh"
1116  try:
1117    with open(os.path.join(input_dir, "BOOT", "RAMDISK", "init.rc")) as f:
1118      for line in f:
1119        m = re.match("^service flash_recovery /system/(\S+)\s*$", line)
1120        if m:
1121          sh_location = m.group(1)
1122          print "putting script in", sh_location
1123          break
1124  except (OSError, IOError), e:
1125    print "failed to read init.rc: %s" % (e,)
1126
1127  output_sink(sh_location, sh)
1128