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