Add firmware image packing tool

utility/pack_image

Index: utility/pack_image
diff --git a/utility/pack_image b/utility/pack_image
new file mode 100755
index 0000000000000000000000000000000000000000..2c9d6ce4cf2862dcaf8a43b92686addd1144d730
--- /dev/null
+++ b/utility/pack_image
@@ -0,0 +1,243 @@
+#!/usr/bin/env python2.6
+
+# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import mmap

[Hung-Te, 2011/01/04 08:39:03]

Any reason to use mmap? I think this script does not really need mmap - file i/o
+ python native array would work.

[Che-Liang Chiou, 2011/01/04 09:38:25]

Because it's supposedly faster.  Do you have strong reason against mmap?

[Hung-Te, 2011/01/04 10:05:23]

  I think performance is not really critical for this script. mmap is more
implementation dependent and may have potential compatibility issues (at least,
forced you to make a hack by touching the file), so I'd prefer general file i/o
which is also easier for maintenance.

[Che-Liang Chiou, 2011/01/05 08:01:55]

Done.

+import os
+import re
+import struct
+import subprocess
+import sys
+import tempfile
+
+FMAP_SIGNATURE = "__FMAP__"

[Hung-Te, 2011/01/04 08:39:03]

The fmap has an python implementation, supporting both encode and decode. We
should reuse the official source as much as possible. In other words, "import
fmap" would be better instead of merging fmap python code in this script.

If you think the official one cannot support what you want to do, please send a
patch to David Hendricks.

[Che-Liang Chiou, 2011/01/04 09:38:25]

I agree, but it should be done in a separate CL because there might be objection
(though unlikely) on pulling David's fmap implementation into chromium os code
base.  I will add a TODO here and rewrite this after we can and does pull in the
fmap.

[Hung-Te, 2011/01/04 10:05:23]

   We already have that in autotest common lib folder.
   Unless you need lots of changes to the official fmap, otherwise I think we
should stop creating more fmap implementation variants.

   Maybe it's a good time to add flashmap as "host-depends" so that developers
can use it (both python and C library).
   

[Che-Liang Chiou, 2011/01/05 08:01:55]

Done.

+
+FMAP_AREA_STATIC = 1 << 0
+FMAP_AREA_COMPRESSED = 1 << 1
+FMAP_AREA_RO = 1 << 2
+
+FMAP_STRLEN = 32
+FMAP_HEADER_FORMAT = "<8sBBQI%dsH" % FMAP_STRLEN
+FMAP_AREA_FORMAT = "<II%dsH" % FMAP_STRLEN
+
+FMAP_HEADER_NAMES = (
+    'signature',
+    'ver_major',
+    'ver_minor',
+    'base',
+    'size',
+    'name',
+    'nareas',
+)
+
+FMAP_AREA_NAMES = (
+    'offset',
+    'size',
+    'name',
+    'flags',
+)
+
+RE_ASSIGNMENT = re.compile(r'^(\w+)=(.*)$')
+RE_HEXADECIMAL = re.compile(r'^0x[0-9A-Fa-f]+$')
+RE_DECIMAL = re.compile(r'^[1-9][0-9]*$')

[Hung-Te, 2011/01/04 08:39:03]

int(a, 0) prevents the use of RE_HEXADECIMAL and RE_DECIMAL.

[Che-Liang Chiou, 2011/01/04 09:38:25]

Done, but this would add one more try-catch block though.

+
+VERBOSE = False
+

[Hung-Te, 2011/01/04 08:39:03]

Please make 2 blank lines for top level definitions (and for all over this
file).

[Che-Liang Chiou, 2011/01/04 09:38:25]

Done.

+class ConfigError(Exception):
+  pass
+
+class PackError(Exception):
+  pass
+
+class Entry(dict):
+  @staticmethod
+  def _check_fields(kwargs, fields):
+    for f in fields:
+      if f not in kwargs:
+        raise ConfigError('Entry: missing required field: %s' % f)
+
+  def __init__(self, **kwargs):
+    Entry._check_fields(kwargs, ('offset', 'length', 'name'))
+    super(Entry, self).__init__(kwargs)
+
+  def __getattr__(self, name):
+    return self[name]
+
+  def is_overlapped(self, entry):
+    return entry.offset <= self.offset < entry.offset + entry.length or \
+        self.offset <= entry.offset < self.offset + self.length
+
+  def pack(self, fw_image, entries):
+    raise PackError('class Entry does not implement pack()')
+
+class EntryFmap(Entry):
+  def __init__(self, **kwargs):
+    Entry._check_fields(kwargs, ('ver_major', 'ver_minor', 'base', 'size'))
+    super(EntryFmap, self).__init__(**kwargs)
+
+  def pack(self, fw_image, entries):
+    areas = [e for e in entries
+        if isinstance(e, EntryBlob) or isinstance(e, EntryKeyBlock)]
+    areas.sort(key=lambda a: a.offset)
+
+    offset = 0
+
+    # pack header
+    values = []
+    for name in FMAP_HEADER_NAMES:
+      if name == 'nareas':
+        v = len(areas)
+      elif name == 'signature':
+        v = FMAP_SIGNATURE
+      else:
+        v = self[name]
+      values.append(v)
+    offset = self._append(fw_image, offset,
+        struct.pack(FMAP_HEADER_FORMAT, *values))
+
+    # pack areas
+    for a in areas:
+      values = [a.length if name == 'size' else a[name]
+          for name in FMAP_AREA_NAMES]
+      offset = self._append(fw_image, offset,
+          struct.pack(FMAP_AREA_FORMAT, *values))
+
+    if offset > self.length:
+        raise PackError('fmap too large: %d > %d' % (offset, self.length))
+
+  def _append(self, fw_image, offset, blob):
+    size = len(blob)
+    fw_image[self.offset+offset:self.offset+offset+size] = blob
+    return offset + size
+
+class EntryBlob(Entry):
+  def __init__(self, **kwargs):
+    Entry._check_fields(kwargs, ('flags', 'path'))
+    super(EntryBlob, self).__init__(**kwargs)
+
+  def pack(self, fw_image, entries):
+    size = os.stat(self.path).st_size
+    if size > 0:
+      size = min(size, self.length)
+    else:
+      size = self.length
+    with open(self.path, 'r+b') as f:
+      fvimg = mmap.mmap(f.fileno(), size, mmap.MAP_PRIVATE, mmap.PROT_READ)
+      fw_image.seek(self.offset)
+      fw_image.write(fvimg[0:size])
+      fvimg.close()
+
+class EntryKeyBlock(Entry):
+  def __init__(self, **kwargs):
+    Entry._check_fields(kwargs,
+        ('flags', 'keyblock', 'signprivate', 'version', 'fv', 'kernelkey'))
+    super(EntryKeyBlock, self).__init__(**kwargs)
+
+  def pack(self, fw_image, entries):
+    fd, path = tempfile.mkstemp()
+    try:
+      args = [
+          'vbutil_firmware',
+          '--vblock', path,
+          '--keyblock', self.keyblock,
+          '--signprivate', self.signprivate,
+          '--version', '%d' % self.version,
+          '--fv', self.fv,
+          '--kernelkey', self.kernelkey,
+      ]
+      if VERBOSE:
+        stdout = sys.stdout
+        stderr = sys.stderr
+      else:
+        stdout = None
+        stderr = None
+      info('run: %s' % ' '.join(args))
+      proc = subprocess.Popen(args, stdout=stdout, stderr=stderr)
+      proc.wait()
+      if proc.returncode != 0:
+        raise PackError('cannot make key block: vbutil_firmware returns %d' %
+			proc.returncode)
+
+      img = mmap.mmap(fd, 0, mmap.MAP_PRIVATE, mmap.PROT_READ)
+      size = img.size()
+      if size > self.length:
+        img.close()
+        raise PackError('key block too large: %d > %d' % (size, self.length))
+      fw_image.seek(self.offset)
+      fw_image.write(img)
+      img.close()
+    finally:
+      os.unlink(path)
+
+def parse_assignment(stmt):
+  m = RE_ASSIGNMENT.match(stmt)
+  if m is None:
+    raise ConfigError('illegal statement: %s' % repr(stmt))
+  return (m.group(1), parse_value(m.group(2)))
+
+def parse_value(expr):
+  if expr[0] == expr[-1] == '"' or expr[0] == expr[-1] == "'":

[Hung-Te, 2011/01/04 08:39:03]

To prevent confusion and help readbility, please add the ()'s  for (expr1) or
(expr2).

Also, startswith and endswith are more readable than a[0] == a[-1] == 'blah'.

[Che-Liang Chiou, 2011/01/04 09:38:25]

Done.

+    expr = expr[1:-1]
+  if RE_HEXADECIMAL.match(expr):
+    return int(expr, 16)
+  elif RE_DECIMAL.match(expr):
+    return int(expr)
+  else:
+    return expr # if not a number, interpret as string literals

[Hung-Te, 2011/01/04 08:39:03]

I'd prefer a more strict syntax, ex: every string must be quoted, otherwise it's
a hex/decimal.

[Che-Liang Chiou, 2011/01/04 09:38:25]

That's not a good idea.  shell will unquote command line arguments before firing
processes.
$ echo -e 'import sys\nprint sys.argv' | python - "hello" "world"
['-', 'hello', 'world']

Or do you suggest we double-quote strings on command line?  I don't prefer this
way though.
$ echo -e 'import sys\nprint sys.argv' | python - "'hello'" "'world'"

[Hung-Te, 2011/01/04 10:05:23]

If you have only 3 parameters, maybe we should consider use getopt and list them
as --opt, instead of the format A=B.

+
+def pack_image(entries, output_path, image_size):
+  entries = sorted(entries, key=lambda e: e.offset)
+  for e1, e2 in zip(entries, entries[1:]):
+    if e1.is_overlapped(e2):
+      raise PackError('overlapped entries: [%08x:%08x], [%08x:%08x]' %
+          (e1.offset, e1.offset + e1.length, e2.offset, e2.offset + e2.length))
+
+  # create image so that mmap will succeed

[Hung-Te, 2011/01/04 08:39:03]

If you use simply buffer processing without mmap, there's no need for such trick
to touch file.

[Che-Liang Chiou, 2011/01/05 08:01:55]

Done.

+  with open(output_path, 'w+b') as f:
+    f.write('\0')
+
+  with open(output_path, 'rw+b') as f:
+    fw_image = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_WRITE)
+    fw_image.resize(image_size)
+    for entry in entries:
+      entry.pack(fw_image, entries)
+    fw_image.flush()
+    fw_image.close()
+
+def info(msg):
+  if VERBOSE:
+    print >>sys.stderr, 'INFO: %s' % msg
+
+def main():
+  global VERBOSE
+
+  if len(sys.argv) < 2:
+    print 'Usage: %s [-v] CONFIG_FILE [NAME=VALUE...]' % sys.argv[0]
+    sys.exit(1)
+
+  if sys.argv[1] == '-v':
+    VERBOSE = True
+    argv = sys.argv[0:1] + sys.argv[2:]
+  else:
+    argv = sys.argv
+
+  if len(argv) > 2:
+    env = dict(parse_assignment(stmt) for stmt in argv[2:])
+  else:
+    env = {}
+
+  execfile(argv[1], globals(), env)
+
+  for varname in ('ENTRIES', 'OUTPUT', 'SIZE'):
+    if varname not in env:
+      raise ConfigError('undefined variable: %s' % varname)
+    info('%s = %s' % (varname, repr(env[varname])))
+
+  pack_image(env['ENTRIES'], env['OUTPUT'], env['SIZE'])
+
+  sys.exit(0)
+
+if __name__ == '__main__':
+  main()