common.py revision fdd8e69c42e66fb70384bcaca1747f504f2c021c 
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 errno
16import getopt
17import getpass
18import os
19import re
20import shutil
21import subprocess
22import sys
23import tempfile
24import zipfile
25
26# missing in Python 2.4 and before
27if not hasattr(os, "SEEK_SET"):
28  os.SEEK_SET = 0
29
30class Options(object): pass
31OPTIONS = Options()
32OPTIONS.search_path = "out/host/linux-x86"
33OPTIONS.max_image_size = {}
34OPTIONS.verbose = False
35OPTIONS.tempfiles = []
36
37
38class ExternalError(RuntimeError): pass
39
40
41def Run(args, **kwargs):
42  """Create and return a subprocess.Popen object, printing the command
43  line on the terminal if -v was specified."""
44  if OPTIONS.verbose:
45    print "  running: ", " ".join(args)
46  return subprocess.Popen(args, **kwargs)
47
48
49def LoadMaxSizes():
50  """Load the maximum allowable images sizes from the input
51  target_files size."""
52  OPTIONS.max_image_size = {}
53  try:
54    for line in open(os.path.join(OPTIONS.input_tmp, "META", "imagesizes.txt")):
55      image, size = line.split()
56      size = int(size)
57      OPTIONS.max_image_size[image + ".img"] = size
58  except IOError, e:
59    if e.errno == errno.ENOENT:
60      pass
61
62
63def BuildAndAddBootableImage(sourcedir, targetname, output_zip):
64  """Take a kernel, cmdline, and ramdisk directory from the input (in
65  'sourcedir'), and turn them into a boot image.  Put the boot image
66  into the output zip file under the name 'targetname'.  Returns
67  targetname on success or None on failure (if sourcedir does not
68  appear to contain files for the requested image)."""
69
70  print "creating %s..." % (targetname,)
71
72  img = BuildBootableImage(sourcedir)
73  if img is None:
74    return None
75
76  CheckSize(img, targetname)
77  ZipWriteStr(output_zip, targetname, img)
78  return targetname
79
80def BuildBootableImage(sourcedir):
81  """Take a kernel, cmdline, and ramdisk directory from the input (in
82  'sourcedir'), and turn them into a boot image.  Return the image
83  data, or None if sourcedir does not appear to contains files for
84  building the requested image."""
85
86  if (not os.access(os.path.join(sourcedir, "RAMDISK"), os.F_OK) or
87      not os.access(os.path.join(sourcedir, "kernel"), os.F_OK)):
88    return None
89
90  ramdisk_img = tempfile.NamedTemporaryFile()
91  img = tempfile.NamedTemporaryFile()
92
93  p1 = Run(["mkbootfs", os.path.join(sourcedir, "RAMDISK")],
94           stdout=subprocess.PIPE)
95  p2 = Run(["minigzip"],
96           stdin=p1.stdout, stdout=ramdisk_img.file.fileno())
97
98  p2.wait()
99  p1.wait()
100  assert p1.returncode == 0, "mkbootfs of %s ramdisk failed" % (targetname,)
101  assert p2.returncode == 0, "minigzip of %s ramdisk failed" % (targetname,)
102
103  cmd = ["mkbootimg", "--kernel", os.path.join(sourcedir, "kernel")]
104
105  fn = os.path.join(sourcedir, "cmdline")
106  if os.access(fn, os.F_OK):
107    cmd.append("--cmdline")
108    cmd.append(open(fn).read().rstrip("\n"))
109
110  fn = os.path.join(sourcedir, "base")
111  if os.access(fn, os.F_OK):
112    cmd.append("--base")
113    cmd.append(open(fn).read().rstrip("\n"))
114
115  cmd.extend(["--ramdisk", ramdisk_img.name,
116              "--output", img.name])
117
118  p = Run(cmd, stdout=subprocess.PIPE)
119  p.communicate()
120  assert p.returncode == 0, "mkbootimg of %s image failed" % (
121      os.path.basename(sourcedir),)
122
123  img.seek(os.SEEK_SET, 0)
124  data = img.read()
125
126  ramdisk_img.close()
127  img.close()
128
129  return data
130
131
132def AddRecovery(output_zip):
133  BuildAndAddBootableImage(os.path.join(OPTIONS.input_tmp, "RECOVERY"),
134                           "recovery.img", output_zip)
135
136def AddBoot(output_zip):
137  BuildAndAddBootableImage(os.path.join(OPTIONS.input_tmp, "BOOT"),
138                           "boot.img", output_zip)
139
140def UnzipTemp(filename):
141  """Unzip the given archive into a temporary directory and return the name."""
142
143  tmp = tempfile.mkdtemp(prefix="targetfiles-")
144  OPTIONS.tempfiles.append(tmp)
145  p = Run(["unzip", "-q", filename, "-d", tmp], stdout=subprocess.PIPE)
146  p.communicate()
147  if p.returncode != 0:
148    raise ExternalError("failed to unzip input target-files \"%s\"" %
149                        (filename,))
150  return tmp
151
152
153def GetKeyPasswords(keylist):
154  """Given a list of keys, prompt the user to enter passwords for
155  those which require them.  Return a {key: password} dict.  password
156  will be None if the key has no password."""
157
158  no_passwords = []
159  need_passwords = []
160  devnull = open("/dev/null", "w+b")
161  for k in sorted(keylist):
162    # An empty-string key is used to mean don't re-sign this package.
163    # Obviously we don't need a password for this non-key.
164    if not k:
165      no_passwords.append(k)
166      continue
167
168    p = Run(["openssl", "pkcs8", "-in", k+".pk8",
169             "-inform", "DER", "-nocrypt"],
170            stdin=devnull.fileno(),
171            stdout=devnull.fileno(),
172            stderr=subprocess.STDOUT)
173    p.communicate()
174    if p.returncode == 0:
175      no_passwords.append(k)
176    else:
177      need_passwords.append(k)
178  devnull.close()
179
180  key_passwords = PasswordManager().GetPasswords(need_passwords)
181  key_passwords.update(dict.fromkeys(no_passwords, None))
182  return key_passwords
183
184
185def SignFile(input_name, output_name, key, password, align=None):
186  """Sign the input_name zip/jar/apk, producing output_name.  Use the
187  given key and password (the latter may be None if the key does not
188  have a password.
189
190  If align is an integer > 1, zipalign is run to align stored files in
191  the output zip on 'align'-byte boundaries.
192  """
193  if align == 0 or align == 1:
194    align = None
195
196  if align:
197    temp = tempfile.NamedTemporaryFile()
198    sign_name = temp.name
199  else:
200    sign_name = output_name
201
202  p = Run(["java", "-jar",
203           os.path.join(OPTIONS.search_path, "framework", "signapk.jar"),
204           key + ".x509.pem",
205           key + ".pk8",
206           input_name, sign_name],
207          stdin=subprocess.PIPE,
208          stdout=subprocess.PIPE)
209  if password is not None:
210    password += "\n"
211  p.communicate(password)
212  if p.returncode != 0:
213    raise ExternalError("signapk.jar failed: return code %s" % (p.returncode,))
214
215  if align:
216    p = Run(["zipalign", "-f", str(align), sign_name, output_name])
217    p.communicate()
218    if p.returncode != 0:
219      raise ExternalError("zipalign failed: return code %s" % (p.returncode,))
220    temp.close()
221
222
223def CheckSize(data, target):
224  """Check the data string passed against the max size limit, if
225  any, for the given target.  Raise exception if the data is too big.
226  Print a warning if the data is nearing the maximum size."""
227  limit = OPTIONS.max_image_size.get(target, None)
228  if limit is None: return
229
230  size = len(data)
231  pct = float(size) * 100.0 / limit
232  msg = "%s size (%d) is %.2f%% of limit (%d)" % (target, size, pct, limit)
233  if pct >= 99.0:
234    raise ExternalError(msg)
235  elif pct >= 95.0:
236    print
237    print "  WARNING: ", msg
238    print
239  elif OPTIONS.verbose:
240    print "  ", msg
241
242
243COMMON_DOCSTRING = """
244  -p  (--path)  <dir>
245      Prepend <dir>/bin to the list of places to search for binaries
246      run by this script, and expect to find jars in <dir>/framework.
247
248  -v  (--verbose)
249      Show command lines being executed.
250
251  -h  (--help)
252      Display this usage message and exit.
253"""
254
255def Usage(docstring):
256  print docstring.rstrip("\n")
257  print COMMON_DOCSTRING
258
259
260def ParseOptions(argv,
261                 docstring,
262                 extra_opts="", extra_long_opts=(),
263                 extra_option_handler=None):
264  """Parse the options in argv and return any arguments that aren't
265  flags.  docstring is the calling module's docstring, to be displayed
266  for errors and -h.  extra_opts and extra_long_opts are for flags
267  defined by the caller, which are processed by passing them to
268  extra_option_handler."""
269
270  try:
271    opts, args = getopt.getopt(
272        argv, "hvp:" + extra_opts,
273        ["help", "verbose", "path="] + list(extra_long_opts))
274  except getopt.GetoptError, err:
275    Usage(docstring)
276    print "**", str(err), "**"
277    sys.exit(2)
278
279  path_specified = False
280
281  for o, a in opts:
282    if o in ("-h", "--help"):
283      Usage(docstring)
284      sys.exit()
285    elif o in ("-v", "--verbose"):
286      OPTIONS.verbose = True
287    elif o in ("-p", "--path"):
288      OPTIONS.search_path = a
289    else:
290      if extra_option_handler is None or not extra_option_handler(o, a):
291        assert False, "unknown option \"%s\"" % (o,)
292
293  os.environ["PATH"] = (os.path.join(OPTIONS.search_path, "bin") +
294                        os.pathsep + os.environ["PATH"])
295
296  return args
297
298
299def Cleanup():
300  for i in OPTIONS.tempfiles:
301    if os.path.isdir(i):
302      shutil.rmtree(i)
303    else:
304      os.remove(i)
305
306
307class PasswordManager(object):
308  def __init__(self):
309    self.editor = os.getenv("EDITOR", None)
310    self.pwfile = os.getenv("ANDROID_PW_FILE", None)
311
312  def GetPasswords(self, items):
313    """Get passwords corresponding to each string in 'items',
314    returning a dict.  (The dict may have keys in addition to the
315    values in 'items'.)
316
317    Uses the passwords in $ANDROID_PW_FILE if available, letting the
318    user edit that file to add more needed passwords.  If no editor is
319    available, or $ANDROID_PW_FILE isn't define, prompts the user
320    interactively in the ordinary way.
321    """
322
323    current = self.ReadFile()
324
325    first = True
326    while True:
327      missing = []
328      for i in items:
329        if i not in current or not current[i]:
330          missing.append(i)
331      # Are all the passwords already in the file?
332      if not missing: return current
333
334      for i in missing:
335        current[i] = ""
336
337      if not first:
338        print "key file %s still missing some passwords." % (self.pwfile,)
339        answer = raw_input("try to edit again? [y]> ").strip()
340        if answer and answer[0] not in 'yY':
341          raise RuntimeError("key passwords unavailable")
342      first = False
343
344      current = self.UpdateAndReadFile(current)
345
346  def PromptResult(self, current):
347    """Prompt the user to enter a value (password) for each key in
348    'current' whose value is fales.  Returns a new dict with all the
349    values.
350    """
351    result = {}
352    for k, v in sorted(current.iteritems()):
353      if v:
354        result[k] = v
355      else:
356        while True:
357          result[k] = getpass.getpass("Enter password for %s key> "
358                                      % (k,)).strip()
359          if result[k]: break
360    return result
361
362  def UpdateAndReadFile(self, current):
363    if not self.editor or not self.pwfile:
364      return self.PromptResult(current)
365
366    f = open(self.pwfile, "w")
367    os.chmod(self.pwfile, 0600)
368    f.write("# Enter key passwords between the [[[ ]]] brackets.\n")
369    f.write("# (Additional spaces are harmless.)\n\n")
370
371    first_line = None
372    sorted = [(not v, k, v) for (k, v) in current.iteritems()]
373    sorted.sort()
374    for i, (_, k, v) in enumerate(sorted):
375      f.write("[[[  %s  ]]] %s\n" % (v, k))
376      if not v and first_line is None:
377        # position cursor on first line with no password.
378        first_line = i + 4
379    f.close()
380
381    p = Run([self.editor, "+%d" % (first_line,), self.pwfile])
382    _, _ = p.communicate()
383
384    return self.ReadFile()
385
386  def ReadFile(self):
387    result = {}
388    if self.pwfile is None: return result
389    try:
390      f = open(self.pwfile, "r")
391      for line in f:
392        line = line.strip()
393        if not line or line[0] == '#': continue
394        m = re.match(r"^\[\[\[\s*(.*?)\s*\]\]\]\s*(\S+)$", line)
395        if not m:
396          print "failed to parse password file: ", line
397        else:
398          result[m.group(2)] = m.group(1)
399      f.close()
400    except IOError, e:
401      if e.errno != errno.ENOENT:
402        print "error reading password file: ", str(e)
403    return result
404
405
406def ZipWriteStr(zip, filename, data, perms=0644):
407  # use a fixed timestamp so the output is repeatable.
408  zinfo = zipfile.ZipInfo(filename=filename,
409                          date_time=(2009, 1, 1, 0, 0, 0))
410  zinfo.compress_type = zip.compression
411  zinfo.external_attr = perms << 16
412  zip.writestr(zinfo, data)
413