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