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