Upstreaming changes from https://chromium.googlesource.com/chromiumos/third_party/flashmap/

.gitignore

@@ -2,6 +2,7 @@
 *.cmd
 *.gcda
 *.gcno
+*.pyc
 .config
 .config.old
 scripts/basic/docproc

OWNERS

@@ -1,2 +1,2 @@
-dhendrix@chromium.org
 dlaurie@chromium.org
+dlaurie@google.com

PRESUBMIT.cfg

@@ -0,0 +1,10 @@
+# This project uses BSD/GPLv2 for the license and tabs for indentation.
+
+[Hook Scripts]
+fmap_unittest py2 = python2 ./fmap_unittest.py
+fmap_unittest py3 = python3 ./fmap_unittest.py
+cros lint = cros lint ${PRESUBMIT_FILES}
+
+[Hook Overrides]
+cros_license_check: false
+tab_check: false

fmap.py

@@ -1,4 +1,5 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
 #
 # Copyright 2010, Google Inc.
 # All rights reserved.
@@ -33,9 +34,7 @@
 # GNU General Public License ("GPL") version 2 as published by the Free
 # Software Foundation.
 
-"""
-This module provides basic encode and decode functionality to the flashrom
-memory map (FMAP) structure.
+"""Basic encode/decode functionality of flashrom memory map (FMAP) structures.
 
 Usage:
   (decode)
@@ -51,18 +50,29 @@
   tuple of decoded area flags.
 """
 
+from __future__ import print_function
+
+import argparse
+import copy
+import logging
+import pprint
 import struct
 import sys
 
+
 # constants imported from lib/fmap.h
-FMAP_SIGNATURE = "__FMAP__"
+FMAP_SIGNATURE = b'__FMAP__'
 FMAP_VER_MAJOR = 1
-FMAP_VER_MINOR = 0
+FMAP_VER_MINOR_MIN = 0
+FMAP_VER_MINOR_MAX = 1
 FMAP_STRLEN = 32
+FMAP_SEARCH_STRIDE = 4
 
 FMAP_FLAGS = {
     'FMAP_AREA_STATIC': 1 << 0,
     'FMAP_AREA_COMPRESSED': 1 << 1,
+    'FMAP_AREA_RO': 1 << 2,
+    'FMAP_AREA_PRESERVE': 1 << 3,
 }
 
 FMAP_HEADER_NAMES = (
@@ -82,13 +92,14 @@
     'flags',
 )
 
+
 # format string
-FMAP_HEADER_FORMAT = "<8sBBQI%dsH" % (FMAP_STRLEN)
-FMAP_AREA_FORMAT = "<II%dsH" % (FMAP_STRLEN)
+FMAP_HEADER_FORMAT = '<8sBBQI%dsH' % (FMAP_STRLEN)
+FMAP_AREA_FORMAT = '<II%dsH' % (FMAP_STRLEN)
 
 
 def _fmap_decode_header(blob, offset):
-  """ (internal) Decodes a FMAP header from blob by offset"""
+  """(internal) Decodes a FMAP header from blob by offset"""
   header = {}
   for (name, value) in zip(FMAP_HEADER_NAMES,
                            struct.unpack_from(FMAP_HEADER_FORMAT,
@@ -98,71 +109,155 @@ def _fmap_decode_header(blob, offset):
 
   if header['signature'] != FMAP_SIGNATURE:
     raise struct.error('Invalid signature')
-  if header['ver_major'] != FMAP_VER_MAJOR or \
-     header['ver_minor'] != FMAP_VER_MINOR:
+  if (header['ver_major'] != FMAP_VER_MAJOR or
+      header['ver_minor'] < FMAP_VER_MINOR_MIN or
+      header['ver_minor'] > FMAP_VER_MINOR_MAX):
     raise struct.error('Incompatible version')
 
   # convert null-terminated names
-  header['name'] = header['name'].strip(chr(0))
+  header['name'] = header['name'].strip(b'\x00')
+
+  # In Python 2, binary==string, so we don't need to convert.
+  if sys.version_info.major >= 3:
+    # Do the decode after verifying it to avoid decode errors due to corruption.
+    for name in FMAP_HEADER_NAMES:
+      if hasattr(header[name], 'decode'):
+        header[name] = header[name].decode('utf-8')
+
   return (header, struct.calcsize(FMAP_HEADER_FORMAT))
 
 
 def _fmap_decode_area(blob, offset):
-  """ (internal) Decodes a FMAP area record from blob by offset """
+  """(internal) Decodes a FMAP area record from blob by offset"""
   area = {}
   for (name, value) in zip(FMAP_AREA_NAMES,
                            struct.unpack_from(FMAP_AREA_FORMAT, blob, offset)):
     area[name] = value
   # convert null-terminated names
-  area['name'] = area['name'].strip(chr(0))
+  area['name'] = area['name'].strip(b'\x00')
   # add a (readonly) readable FLAGS
   area['FLAGS'] = _fmap_decode_area_flags(area['flags'])
+
+  # In Python 2, binary==string, so we don't need to convert.
+  if sys.version_info.major >= 3:
+    for name in FMAP_AREA_NAMES:
+      if hasattr(area[name], 'decode'):
+        area[name] = area[name].decode('utf-8')
+
   return (area, struct.calcsize(FMAP_AREA_FORMAT))
 
 
 def _fmap_decode_area_flags(area_flags):
-  """ (internal) Decodes a FMAP flags property """
-  return tuple([name for name in FMAP_FLAGS if area_flags & FMAP_FLAGS[name]])
+  """(internal) Decodes a FMAP flags property"""
+  # Since FMAP_FLAGS is a dict with arbitrary ordering, sort the list so the
+  # output is stable.  Also sorting is nicer for humans.
+  return tuple(sorted(x for x in FMAP_FLAGS if area_flags & FMAP_FLAGS[x]))
 
 
-def fmap_decode(blob, offset=None):
-  """ Decodes a blob to FMAP dictionary object.
+def _fmap_check_name(fmap, name):
+  """Checks if the FMAP structure has correct name.
 
-  Arguments:
+  Args:
+    fmap: A decoded FMAP structure.
+    name: A string to specify expected FMAP name.
+
+  Raises:
+    struct.error if the name does not match.
+  """
+  if fmap['name'] != name:
+    raise struct.error('Incorrect FMAP (found: "%s", expected: "%s")' %
+                       (fmap['name'], name))
+
+
+def _fmap_search_header(blob, fmap_name=None):
+  """Searches FMAP headers in given blob.
+
+  Uses same logic from vboot_reference/host/lib/fmap.c.
+
+  Args:
+    blob: A string containing FMAP data.
+    fmap_name: A string to specify target FMAP name.
+
+  Returns:
+    A tuple of (fmap, size, offset).
+  """
+  lim = len(blob) - struct.calcsize(FMAP_HEADER_FORMAT)
+  align = FMAP_SEARCH_STRIDE
+
+  # Search large alignments before small ones to find "right" FMAP.
+  while align <= lim:
+    align *= 2
+
+  while align >= FMAP_SEARCH_STRIDE:
+    for offset in range(align, lim + 1, align * 2):
+      if not blob.startswith(FMAP_SIGNATURE, offset):
+        continue
+      try:
+        (fmap, size) = _fmap_decode_header(blob, offset)
+        if fmap_name is not None:
+          _fmap_check_name(fmap, fmap_name)
+        return (fmap, size, offset)
+      except struct.error as e:
+        # Search for next FMAP candidate.
+        logging.debug('Continue searching FMAP due to exception %r', e)
+    align //= 2
+  raise struct.error('No valid FMAP signatures.')
+
+
+def fmap_decode(blob, offset=None, fmap_name=None):
+  """Decodes a blob to FMAP dictionary object.
+
+  Args:
     blob: a binary data containing FMAP structure.
     offset: starting offset of FMAP. When omitted, fmap_decode will search in
             the blob.
+    fmap_name: A string to specify target FMAP name.
   """
   fmap = {}
-  if offset == None:
-    # try search magic in fmap
-    offset = blob.find(FMAP_SIGNATURE)
-  (fmap, size) = _fmap_decode_header(blob, offset)
+
+  if offset is None:
+    (fmap, size, offset) = _fmap_search_header(blob, fmap_name)
+  else:
+    (fmap, size) = _fmap_decode_header(blob, offset)
+    if fmap_name is not None:
+      _fmap_check_name(fmap, fmap_name)
   fmap['areas'] = []
   offset = offset + size
-  for i in range(fmap['nareas']):
+  for _ in range(fmap['nareas']):
     (area, size) = _fmap_decode_area(blob, offset)
     offset = offset + size
     fmap['areas'].append(area)
   return fmap
 
 
 def _fmap_encode_header(obj):
-  """ (internal) Encodes a FMAP header """
+  """(internal) Encodes a FMAP header"""
+  # Convert strings to bytes.
+  obj = copy.deepcopy(obj)
+  for name in FMAP_HEADER_NAMES:
+    if hasattr(obj[name], 'encode'):
+      obj[name] = obj[name].encode('utf-8')
+
   values = [obj[name] for name in FMAP_HEADER_NAMES]
   return struct.pack(FMAP_HEADER_FORMAT, *values)
 
 
 def _fmap_encode_area(obj):
-  """ (internal) Encodes a FMAP area entry """
+  """(internal) Encodes a FMAP area entry"""
+  # Convert strings to bytes.
+  obj = copy.deepcopy(obj)
+  for name in FMAP_AREA_NAMES:
+    if hasattr(obj[name], 'encode'):
+      obj[name] = obj[name].encode('utf-8')
+
   values = [obj[name] for name in FMAP_AREA_NAMES]
   return struct.pack(FMAP_AREA_FORMAT, *values)
 
 
 def fmap_encode(obj):
-  """ Encodes a FMAP dictionary object to blob.
+  """Encodes a FMAP dictionary object to blob.
 
-  Arguments
+  Args:
     obj: a FMAP dictionary object.
   """
   # fix up values
@@ -174,13 +269,32 @@ def fmap_encode(obj):
   return blob
 
 
-if __name__ == '__main__':
-  # main entry, do a unit test
-  blob = open('bin/example.bin').read()
+def get_parser():
+  """Return a command line parser."""
+  parser = argparse.ArgumentParser(
+      description=__doc__,
+      formatter_class=argparse.RawTextHelpFormatter)
+  parser.add_argument('file', help='The file to decode & print.')
+  parser.add_argument('--raw', action='store_true',
+                      help='Dump the object output for scripts.')
+  return parser
+
+
+def main(argv):
+  """Decode FMAP from supplied file and print."""
+  parser = get_parser()
+  opts = parser.parse_args(argv)
+
+  if not opts.raw:
+    print('Decoding FMAP from: %s' % opts.file)
+  blob = open(opts.file, 'rb').read()
   obj = fmap_decode(blob)
-  print obj
-  blob2 = fmap_encode(obj)
-  obj2 = fmap_decode(blob2)
-  print obj2
-  assert obj == obj2
+  if opts.raw:
+    print(obj)
+  else:
+    pp = pprint.PrettyPrinter(indent=2)
+    pp.pprint(obj)
+
 
+if __name__ == '__main__':
+  sys.exit(main(sys.argv[1:]))